This is an automated email from the ASF dual-hosted git repository.
kszucs pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push:
new 0894045 ARROW-7670: [Python][Dataset] More ergonomical API
0894045 is described below
commit 089404571bcfb66e079ec70b8984a402e0c08b04
Author: Krisztián Szűcs <[email protected]>
AuthorDate: Sun Jan 26 19:43:31 2020 +0100
ARROW-7670: [Python][Dataset] More ergonomical API
- introduced `ds.field()`, `ds.scalar()`
- overloaded comparison and binary operators
- added `Expression.isin()` and `Expression.is_valid()`
- removed ScannerBuilder in favor of direct `Scanner` instantiation using
keyword arguments
- added convenience methods for Dataset: `to_table()`, `to_batches()` and
`scan()`
- added support for string scalar
- removed `InMemoryScanTask` and `DefaultPartitioning`
Closes #6268 from kszucs/filter-expr and squashes the following commits:
2bcf0554f <Krisztián Szűcs> fix review issues
706f3cc25 <Krisztián Szűcs> add to apidocs
a54a7c806 <Krisztián Szűcs> docstring
0966b8afa <Krisztián Szűcs> flake8
a950c975b <Krisztián Szűcs> flake8
9af49ce29 <Krisztián Szűcs> revert abbreviations
0170362cf <Krisztián Szűcs> cleanup
31cd5c60b <Krisztián Szűcs> remove ScannerBuilder
b0b4cd9b8 <Krisztián Szűcs> fix rename issue
b3438dfb2 <Krisztián Szűcs> cython specific overloading
dcab6b069 <Krisztián Szűcs> rename Expression to Expr
757343c13 <Krisztián Szűcs> string scalar
9c7ee5a03 <Krisztián Szűcs> improve the filter expression ergonomics
Authored-by: Krisztián Szűcs <[email protected]>
Signed-off-by: Krisztián Szűcs <[email protected]>
---
docs/source/python/api/dataset.rst | 3 +
python/pyarrow/_dataset.pyx | 553 +++++++++++++++------------
python/pyarrow/dataset.py | 37 +-
python/pyarrow/includes/libarrow.pxd | 4 +
python/pyarrow/includes/libarrow_dataset.pxd | 16 +-
python/pyarrow/tests/test_dataset.py | 164 ++++----
6 files changed, 450 insertions(+), 327 deletions(-)
diff --git a/docs/source/python/api/dataset.rst
b/docs/source/python/api/dataset.rst
index c2d37d0..e7dea33 100644
--- a/docs/source/python/api/dataset.rst
+++ b/docs/source/python/api/dataset.rst
@@ -36,6 +36,8 @@ Factory functions
source
dataset
partitioning
+ field
+ scalar
Classes
-------
@@ -56,3 +58,4 @@ Classes
Dataset
ScannerBuilder
Scanner
+ Expression
diff --git a/python/pyarrow/_dataset.pyx b/python/pyarrow/_dataset.pyx
index 4cd0cd2..7dbc971 100644
--- a/python/pyarrow/_dataset.pyx
+++ b/python/pyarrow/_dataset.pyx
@@ -21,8 +21,11 @@
from __future__ import absolute_import
+import six
+from cpython.object cimport Py_LT, Py_EQ, Py_GT, Py_LE, Py_NE, Py_GE
from cython.operator cimport dereference as deref
+import pyarrow as pa
from pyarrow.lib cimport *
from pyarrow.includes.libarrow_dataset cimport *
from pyarrow.compat import frombytes, tobytes
@@ -95,9 +98,7 @@ cdef class Partitioning:
cdef Partitioning self
typ = frombytes(sp.get().type_name())
- if typ == 'default':
- self = DefaultPartitioning.__new__(DefaultPartitioning)
- elif typ == 'schema':
+ if typ == 'schema':
self = DirectoryPartitioning.__new__(DirectoryPartitioning)
elif typ == 'hive':
self = HivePartitioning.__new__(HivePartitioning)
@@ -146,21 +147,6 @@ cdef class PartitioningFactory:
return self.wrapped
-cdef class DefaultPartitioning(Partitioning):
-
- cdef:
- CDefaultPartitioning* default_partitioning
-
- def __init__(self):
- cdef shared_ptr[CDefaultPartitioning] partitioning
- partitioning = make_shared[CDefaultPartitioning]()
- self.init(<shared_ptr[CPartitioning]> partitioning)
-
- cdef init(self, const shared_ptr[CPartitioning]& sp):
- Partitioning.init(self, sp)
- self.default_partitioning = <CDefaultPartitioning*> sp.get()
-
-
cdef class DirectoryPartitioning(Partitioning):
"""
A Partitioning based on a specified Schema.
@@ -292,7 +278,7 @@ cdef class HivePartitioning(Partitioning):
cdef class FileSystemFactoryOptions:
"""
- Options for FileSystemFactoryOptions.
+ Influences the discovery of filesystem paths.
Parameters
----------
@@ -492,10 +478,13 @@ cdef class FileSystemSourceFactory(SourceFactory):
Parameters
----------
filesystem : pyarrow.fs.FileSystem
+ Filesystem to discover.
paths_or_selector: pyarrow.fs.Selector or list of path-likes
Either a Selector object or a list of path-like objects.
format : FileFormat
+ Currently only ParquetFileFormat is supported.
options : FileSystemFactoryOptions, optional
+ Various flags influencing the discovery of filesystem paths.
"""
cdef:
@@ -513,7 +502,6 @@ cdef class FileSystemSourceFactory(SourceFactory):
CFileSystemFactoryOptions c_options
c_filesystem = filesystem.unwrap()
-
c_format = format.unwrap()
options = options or FileSystemFactoryOptions()
@@ -589,15 +577,15 @@ cdef class Source:
@property
def partition_expression(self):
"""
- An expression which evaluates to true for all data viewed by this
+ An Expression which evaluates to true for all data viewed by this
Source.
"""
- cdef shared_ptr[CExpression] expression
- expression = self.source.partition_expression()
- if expression.get() == nullptr:
+ cdef shared_ptr[CExpression] expr
+ expr = self.source.partition_expression()
+ if expr.get() == nullptr:
return None
else:
- return Expression.wrap(expression)
+ return Expression.wrap(expr)
cdef class TreeSource(Source):
@@ -630,8 +618,7 @@ cdef class FileSystemSource(Source):
cdef:
CFileSystemSource* filesystem_source
- def __init__(self, Schema schema not None,
- Expression source_partition,
+ def __init__(self, Schema schema not None, Expression root_partition,
FileFormat file_format not None,
FileSystem filesystem not None,
paths_or_selector, partitions):
@@ -640,20 +627,22 @@ cdef class FileSystemSource(Source):
Parameters
----------
schema : Schema
- Schema for resulting Source
- source_partition : Expression
+ The top-level schema of the DataSource.
+ root_partition : Expression
+ The top-level partition of the DataSource.
file_format : FileFormat
+ File format to create fragments from, currently only
+ ParquetFileFormat is supported.
filesystem : FileSystem
- FileSystem which will be explored to discover data files.
+ The filesystem which files are from.
paths_or_selector : Union[FileSelector, List[FileStats]]
- The file stats object can be queried by the
- filesystem.get_target_stats method.
+ List of files/directories to consume.
partitions : List[Expression]
+ Attach aditional partition information for the file paths.
"""
cdef:
- shared_ptr[CExpression] c_source_partition
FileStats stats
- Expression expression
+ Expression expr
vector[CFileStats] c_file_stats
vector[shared_ptr[CExpression]] c_partitions
CResult[shared_ptr[CSource]] result
@@ -661,8 +650,8 @@ cdef class FileSystemSource(Source):
for stats in filesystem.get_target_stats(paths_or_selector):
c_file_stats.push_back(stats.unwrap())
- for expression in partitions:
- c_partitions.push_back(expression.unwrap())
+ for expr in partitions:
+ c_partitions.push_back(expr.unwrap())
if c_file_stats.size() != c_partitions.size():
raise ValueError(
@@ -670,13 +659,12 @@ cdef class FileSystemSource(Source):
'equal to the number of partitions.'
)
- if source_partition is None:
- source_partition = ScalarExpression(True)
- c_source_partition = source_partition.unwrap()
+ if root_partition is None:
+ root_partition = ScalarExpression(True)
result = CFileSystemSource.Make(
pyarrow_unwrap_schema(schema),
- c_source_partition,
+ root_partition.unwrap(),
file_format.unwrap(),
filesystem.unwrap(),
c_file_stats,
@@ -805,30 +793,120 @@ cdef class Dataset:
cdef inline shared_ptr[CDataset] unwrap(self) nogil:
return self.wrapped
- def new_scan(self, MemoryPool memory_pool=None):
+ def scan(self, columns=None, filter=None, MemoryPool memory_pool=None):
+ """Builds a scan operation against the dataset.
+
+ It poduces a stream of ScanTasks which is meant to be a unit of work to
+ be dispatched. The tasks are not executed automatically, the user is
+ responsible to execute and dispatch the individual tasks, so custom
+ local task scheduling can be implemented.
+
+ Parameters
+ ----------
+ columns : list of str, default None
+ List of columns to project. Order and duplicates will be preserved.
+ The columns will be passed down to Sources and corresponding data
+ fragments to avoid loading, copying, and deserializing columns
+ that will not be required further down the compute chain.
+ By default all of the available columns are projected. Raises
+ an exception if any of the referenced column names does not exist
+ in the dataset's Schema.
+ filter : Expression, default None
+ Scan will return only the rows matching the filter.
+ If possible the predicate will be pushed down to exploit the
+ partition information or internal metadata found in the data
+ source, e.g. Parquet statistics. Otherwise filters the loaded
+ RecordBatches before yielding them.
+ memory_pool : MemoryPool, default None
+ For memory allocations, if required. If not specified, uses the
+ default pool.
+
+ Returns
+ -------
+ scan_tasks : iterator of ScanTask
"""
- Begin to build a new Scan operation against this Dataset.
+ scanner = Scanner(self, columns=columns, filter=filter,
+ memory_pool=memory_pool)
+ return scanner.scan()
+
+ def to_batches(self, columns=None, filter=None,
+ MemoryPool memory_pool=None):
+ """Read the dataset as materialized record batches.
+
+ Builds a scan operation against the dataset and sequentially executes
+ the ScanTasks as the returned generator gets consumed.
Parameters
----------
+ columns : list of str, default None
+ List of columns to project. Order and duplicates will be preserved.
+ The columns will be passed down to Sources and corresponding data
+ fragments to avoid loading, copying, and deserializing columns
+ that will not be required further down the compute chain.
+ By default all of the available columns are projected. Raises
+ an exception if any of the referenced column names does not exist
+ in the dataset's Schema.
+ filter : Expression, default None
+ Scan will return only the rows matching the filter.
+ If possible the predicate will be pushed down to exploit the
+ partition information or internal metadata found in the data
+ source, e.g. Parquet statistics. Otherwise filters the loaded
+ RecordBatches before yielding them.
memory_pool : MemoryPool, default None
For memory allocations, if required. If not specified, uses the
default pool.
Returns
-------
- ScannerBuilder
+ record_batches : iterator of RecordBatch
"""
- cdef:
- shared_ptr[CScanContext] context = make_shared[CScanContext]()
- CResult[shared_ptr[CScannerBuilder]] result
- context.get().pool = maybe_unbox_memory_pool(memory_pool)
- result = self.dataset.NewScanWithContext(context)
- return ScannerBuilder.wrap(GetResultValue(result))
+ scanner = Scanner(self, columns=columns, filter=filter,
+ memory_pool=memory_pool)
+ for task in scanner.scan():
+ for batch in task.execute():
+ yield batch
+
+ def to_table(self, columns=None, filter=None, use_threads=True,
+ MemoryPool memory_pool=None):
+ """Read the dataset to an arrow table.
+
+ Note that this method reads all the selected data from the dataset
+ into memory.
+
+ Parameters
+ ----------
+ columns : list of str, default None
+ List of columns to project. Order and duplicates will be preserved.
+ The columns will be passed down to Sources and corresponding data
+ fragments to avoid loading, copying, and deserializing columns
+ that will not be required further down the compute chain.
+ By default all of the available columns are projected. Raises
+ an exception if any of the referenced column names does not exist
+ in the dataset's Schema.
+ filter : Expression, default None
+ Scan will return only the rows matching the filter.
+ If possible the predicate will be pushed down to exploit the
+ partition information or internal metadata found in the data
+ source, e.g. Parquet statistics. Otherwise filters the loaded
+ RecordBatches before yielding them.
+ use_threads : boolean, default True
+ If enabled, then maximum paralellism will be used determined by
+ the number of available CPU cores.
+ memory_pool : MemoryPool, default None
+ For memory allocations, if required. If not specified, uses the
+ default pool.
+
+ Returns
+ -------
+ table : Table instance
+ """
+ scanner = Scanner(self, columns=columns, filter=filter,
+ use_threads=use_threads, memory_pool=memory_pool)
+ return scanner.to_table()
@property
def sources(self):
- """List of the sources"""
+ """List of the data sources"""
cdef vector[shared_ptr[CSource]] sources = self.dataset.sources()
return [Source.wrap(source) for source in sources]
@@ -857,7 +935,7 @@ cdef class ScanTask:
@staticmethod
cdef wrap(shared_ptr[CScanTask]& sp):
- cdef InMemoryScanTask self = InMemoryScanTask.__new__(InMemoryScanTask)
+ cdef ScanTask self = ScanTask.__new__(ScanTask)
self.init(sp)
return self
@@ -878,173 +956,90 @@ cdef class ScanTask:
CRecordBatchIterator iterator
shared_ptr[CRecordBatch] record_batch
- iterator = move(GetResultValue(move(self.task.Execute())))
-
- while True:
- record_batch = GetResultValue(iterator.Next())
- if record_batch.get() == nullptr:
- raise StopIteration()
- else:
- yield pyarrow_wrap_batch(record_batch)
-
+ with nogil:
+ iterator = move(GetResultValue(move(self.task.Execute())))
-cdef class InMemoryScanTask(ScanTask):
- """A trivial ScanTask that yields the RecordBatch of an array."""
+ while True:
+ record_batch = GetResultValue(iterator.Next())
+ if record_batch.get() == nullptr:
+ raise StopIteration()
+ else:
+ with gil:
+ yield pyarrow_wrap_batch(record_batch)
- cdef:
- CInMemoryScanTask* in_memory_task
-
- cdef init(self, shared_ptr[CScanTask]& sp):
- ScanTask.init(self, sp)
- self.in_memory_task = <CInMemoryScanTask*> sp.get()
+cdef class Scanner:
+ """A materialized scan operation with context and options bound.
-cdef class ScannerBuilder:
- """Factory class to construct a Scanner.
-
- It is used to pass information, notably a potential filter expression and a
- subset of columns to materialize.
+ A scanner is the class that glues the scan tasks, data fragments and data
+ sources together.
Parameters
----------
dataset : Dataset
- The dataset to scan.
+ Dataset to scan.
+ columns : list of str, default None
+ List of columns to project. Order and duplicates will be preserved.
+ The columns will be passed down to Sources and corresponding data
+ fragments to avoid loading, copying, and deserializing columns
+ that will not be required further down the compute chain.
+ By default all of the available columns are projected. Raises
+ an exception if any of the referenced column names does not exist
+ in the dataset's Schema.
+ filter : Expression, default None
+ Scan will return only the rows matching the filter.
+ If possible the predicate will be pushed down to exploit the
+ partition information or internal metadata found in the data
+ source, e.g. Parquet statistics. Otherwise filters the loaded
+ RecordBatches before yielding them.
+ use_threads : boolean, default True
+ If enabled, then maximum paralellism will be used determined by
+ the number of available CPU cores.
memory_pool : MemoryPool, default None
For memory allocations, if required. If not specified, uses the
default pool.
"""
cdef:
- shared_ptr[CScannerBuilder] wrapped
- CScannerBuilder* builder
+ shared_ptr[CScanner] wrapped
+ CScanner* scanner
- def __init__(self, Dataset dataset not None, MemoryPool memory_pool=None):
+ def __init__(self, Dataset dataset, list columns=None,
+ Expression filter=None, bint use_threads=True,
+ MemoryPool memory_pool=None):
cdef:
+ shared_ptr[CScanContext] context
shared_ptr[CScannerBuilder] builder
- shared_ptr[CScanContext] context = make_shared[CScanContext]()
- context.get().pool = maybe_unbox_memory_pool(memory_pool)
- builder = make_shared[CScannerBuilder](dataset.unwrap(), context)
- self.init(builder)
-
- cdef void init(self, shared_ptr[CScannerBuilder]& sp):
- self.wrapped = sp
- self.builder = sp.get()
-
- @staticmethod
- cdef wrap(shared_ptr[CScannerBuilder]& sp):
- cdef ScannerBuilder self = ScannerBuilder.__new__(ScannerBuilder)
- self.init(sp)
- return self
-
- cdef inline shared_ptr[CScannerBuilder] unwrap(self) nogil:
- return self.wrapped
-
- def project(self, columns):
- """Set the subset of columns to materialize.
-
- This subset will be passed down to Sources and corresponding
- data fragments. The goal is to avoid loading, copying, and
- deserializing columns that will not be required further down the
- compute chain.
-
- It alters the object in place and returns the object itself enabling
- method chaining. Raises exception if any of the referenced column names
- does not exists in the dataset's Schema.
-
- Parameters
- ----------
- columns : list of str
- List of columns to project. Order and duplicates will be preserved.
+ shared_ptr[CExpression] filter_expression
+ vector[c_string] columns_to_project
- Returns
- -------
- self : ScannerBuilder
- """
- cdef vector[c_string] cols = [tobytes(c) for c in columns]
- check_status(self.builder.Project(cols))
-
- return self
-
- def finish(self):
- """Return the constructed now-immutable Scanner object
-
- Returns
- -------
- Scanner
- """
- return Scanner.wrap(GetResultValue(self.builder.Finish()))
-
- def filter(self, Expression filter_expression not None):
- """Set the filter expression to return only rows matching the filter.
-
- The predicate will be passed down to Sources and corresponding
- data fragments to exploit predicate pushdown if possible using
- partition information or internal metadata, e.g. Parquet statistics.
- Otherwise filters the loaded RecordBatches before yielding them.
-
- It alters the object in place and returns the object itself enabling
- method chaining. Raises exception if any of the referenced column names
- does not exists in the dataset's Schema.
-
- Parameters
- ----------
- filter_expression : Expression
- Boolean expression to filter rows with.
-
- Returns
- -------
- self : ScannerBuilder
- """
- cdef:
- shared_ptr[CExpression] c_casting_filter_expression
+ # create scan context
+ context = make_shared[CScanContext]()
+ context.get().pool = maybe_unbox_memory_pool(memory_pool)
- c_casting_filter_expression = GetResultValue(
- CInsertImplicitCasts(
- deref(filter_expression.unwrap().get()),
- deref(self.builder.schema().get())
- )
+ # create scanner builder
+ builder = GetResultValue(
+ dataset.unwrap().get().NewScanWithContext(context)
)
- check_status(self.builder.Filter(c_casting_filter_expression))
- return self
-
- def use_threads(self, bint value):
- """Set whether the Scanner should make use of the thread pool.
-
- It alters the object in place and returns with the object itself
- enabling method chaining.
-
- Parameters
- ----------
- value : boolean
-
- Returns
- -------
- self : ScannerBuilder
- """
- check_status(self.builder.UseThreads(value))
- return self
-
- @property
- def schema(self):
- return pyarrow_wrap_schema(self.builder.schema())
-
-cdef class Scanner:
- """A materialized scan operation with context and options bound.
-
- Create this using the ScannerBuilder factory class.
-
- A scanner is the class that glues the scan tasks, data fragments and data
- sources together.
- """
-
- cdef:
- shared_ptr[CScanner] wrapped
- CScanner* scanner
+ # set the builder's properties
+ if columns is not None:
+ columns_to_project = [tobytes(c) for c in columns]
+ check_status(builder.get().Project(columns_to_project))
+ if filter is not None:
+ filter_expression = GetResultValue(
+ CInsertImplicitCasts(
+ deref(filter.unwrap().get()),
+ deref(builder.get().schema().get())
+ )
+ )
+ check_status(builder.get().Filter(filter_expression))
+ if use_threads is not None:
+ check_status(builder.get().UseThreads(use_threads))
- def __init__(self):
- raise TypeError('Scanner cannot be initialized directly, use '
- 'ScannerBuilder instead')
+ # instantiate the scanner object
+ scanner = GetResultValue(builder.get().Finish())
+ self.init(scanner)
cdef void init(self, shared_ptr[CScanner]& sp):
self.wrapped = sp
@@ -1097,18 +1092,41 @@ cdef class Scanner:
return pyarrow_wrap_table(GetResultValue(result))
+def _binop(fn, left, right):
+ # cython doesn't support reverse operands like __radd__ just passes the
+ # arguments in the same order as the binary operator called
+
+ if isinstance(left, Expression) and isinstance(right, Expression):
+ pass
+ elif isinstance(left, Expression):
+ try:
+ right = ScalarExpression(right)
+ except TypeError:
+ return NotImplemented
+
+ elif isinstance(right, Expression):
+ try:
+ left = ScalarExpression(left)
+ except TypeError:
+ return NotImplemented
+ else:
+ raise TypeError('Neither left nor right arguments are Expressions')
+
+ return fn(left, right)
+
+
cdef class Expression:
cdef:
shared_ptr[CExpression] wrapped
- CExpression* expression
+ CExpression* expr
def __init__(self):
_forbid_instantiation(self.__class__)
cdef void init(self, const shared_ptr[CExpression]& sp):
self.wrapped = sp
- self.expression = sp.get()
+ self.expr = sp.get()
@staticmethod
cdef wrap(const shared_ptr[CExpression]& sp):
@@ -1143,21 +1161,76 @@ cdef class Expression:
return self.wrapped
def equals(self, Expression other):
- return self.expression.Equals(other.unwrap())
+ return self.expr.Equals(other.unwrap())
def __str__(self):
- return frombytes(self.expression.ToString())
+ return frombytes(self.expr.ToString())
def validate(self, Schema schema not None):
+ """Validate this expression for execution against a schema.
+
+ This will check that all reference fields are present (fields not in
+ the schema will be replaced with null) and all subexpressions are
+ executable. Returns the type to which this expression will evaluate.
+
+ Parameters
+ ----------
+ schema : Schema
+ Schema to execute the expression on.
+
+ Returns
+ -------
+ type : DataType
+ """
cdef:
shared_ptr[CSchema] sp_schema
CResult[shared_ptr[CDataType]] result
sp_schema = pyarrow_unwrap_schema(schema)
- result = self.expression.Validate(deref(sp_schema))
+ result = self.expr.Validate(deref(sp_schema))
return pyarrow_wrap_data_type(GetResultValue(result))
def assume(self, Expression given):
- return Expression.wrap(self.expression.Assume(given.unwrap()))
+ """Simplify to an equivalent Expression given assumed constraints."""
+ return Expression.wrap(self.expr.Assume(given.unwrap()))
+
+ def __invert__(self):
+ return NotExpression(self)
+
+ def __richcmp__(self, other, int op):
+ operator_mapping = {
+ Py_EQ: CompareOperator.Equal,
+ Py_NE: CompareOperator.NotEqual,
+ Py_GT: CompareOperator.Greater,
+ Py_GE: CompareOperator.GreaterEqual,
+ Py_LT: CompareOperator.Less,
+ Py_LE: CompareOperator.LessEqual
+ }
+
+ if not isinstance(other, Expression):
+ try:
+ other = ScalarExpression(other)
+ except TypeError:
+ return NotImplemented
+
+ return ComparisonExpression(operator_mapping[op], self, other)
+
+ def __and__(self, other):
+ return _binop(AndExpression, self, other)
+
+ def __or__(self, other):
+ return _binop(OrExpression, self, other)
+
+ def is_valid(self):
+ """Checks whether the expression is not-null (valid)"""
+ return IsValidExpression(self)
+
+ def cast(self, type, bint safe=True):
+ """Explicitly change the expression's data type"""
+ return CastExpression(self, to=ensure_type(type), safe=safe)
+
+ def isin(self, values):
+ """Checks whether the expression is contained in values"""
+ return InExpression(self, pa.array(values))
cdef class UnaryExpression(Expression):
@@ -1193,7 +1266,7 @@ cdef class ScalarExpression(Expression):
def __init__(self, value):
cdef:
shared_ptr[CScalar] scalar
- shared_ptr[CScalarExpression] expression
+ shared_ptr[CScalarExpression] expr
if isinstance(value, bool):
scalar = MakeScalar(<c_bool>value)
@@ -1201,21 +1274,18 @@ cdef class ScalarExpression(Expression):
scalar = MakeScalar(<double>value)
elif isinstance(value, int):
scalar = MakeScalar(<int64_t>value)
+ elif isinstance(value, six.string_types):
+ scalar = MakeStringScalar(tobytes(value))
else:
raise TypeError('Not yet supported scalar value: {}'.format(value))
- expression.reset(new CScalarExpression(scalar))
- self.init(<shared_ptr[CExpression]> expression)
+ expr.reset(new CScalarExpression(scalar))
+ self.init(<shared_ptr[CExpression]> expr)
cdef void init(self, const shared_ptr[CExpression]& sp):
Expression.init(self, sp)
self.scalar = <CScalarExpression*> sp.get()
- # TODO(kszucs): implement once we have proper Scalar bindings
- # @property
- # def value(self):
- # return pyarrow_wrap_scalar(self.scalar.value())
-
cdef class FieldExpression(Expression):
@@ -1224,9 +1294,9 @@ cdef class FieldExpression(Expression):
def __init__(self, name):
cdef:
c_string field_name = tobytes(name)
- shared_ptr[CExpression] expression
- expression.reset(new CFieldExpression(field_name))
- self.init(expression)
+ shared_ptr[CExpression] expr
+ expr.reset(new CFieldExpression(field_name))
+ self.init(expr)
cdef void init(self, const shared_ptr[CExpression]& sp):
Expression.init(self, sp)
@@ -1249,18 +1319,17 @@ cdef class ComparisonExpression(BinaryExpression):
cdef CComparisonExpression* comparison
- def __init__(self, CompareOperator op,
- Expression left_operand not None,
- Expression right_operand not None):
- cdef shared_ptr[CComparisonExpression] expression
- expression.reset(
+ def __init__(self, CompareOperator op, Expression left not None,
+ Expression right not None):
+ cdef shared_ptr[CComparisonExpression] expr
+ expr.reset(
new CComparisonExpression(
<CCompareOperator>op,
- left_operand.unwrap(),
- right_operand.unwrap()
+ left.unwrap(),
+ right.unwrap()
)
)
- self.init(<shared_ptr[CExpression]> expression)
+ self.init(<shared_ptr[CExpression]> expr)
cdef void init(self, const shared_ptr[CExpression]& sp):
BinaryExpression.init(self, sp)
@@ -1273,9 +1342,9 @@ cdef class ComparisonExpression(BinaryExpression):
cdef class IsValidExpression(UnaryExpression):
def __init__(self, Expression operand not None):
- cdef shared_ptr[CIsValidExpression] expression
- expression = make_shared[CIsValidExpression](operand.unwrap())
- self.init(<shared_ptr[CExpression]> expression)
+ cdef shared_ptr[CIsValidExpression] expr
+ expr = make_shared[CIsValidExpression](operand.unwrap())
+ self.init(<shared_ptr[CExpression]> expr)
cdef class CastExpression(UnaryExpression):
@@ -1286,59 +1355,57 @@ cdef class CastExpression(UnaryExpression):
# term we should expose the CastOptions object
cdef:
CastOptions options
- shared_ptr[CExpression] expression
+ shared_ptr[CExpression] expr
options = CastOptions.safe() if safe else CastOptions.unsafe()
- expression.reset(new CCastExpression(
+ expr.reset(new CCastExpression(
operand.unwrap(),
pyarrow_unwrap_data_type(to),
options.unwrap()
))
- self.init(expression)
+ self.init(expr)
cdef class InExpression(UnaryExpression):
def __init__(self, Expression operand not None, Array haystack not None):
- cdef shared_ptr[CExpression] expression
- expression.reset(
+ cdef shared_ptr[CExpression] expr
+ expr.reset(
new CInExpression(operand.unwrap(), pyarrow_unwrap_array(haystack))
)
- self.init(expression)
+ self.init(expr)
cdef class NotExpression(UnaryExpression):
def __init__(self, Expression operand not None):
- cdef shared_ptr[CNotExpression] expression
- expression = MakeNotExpression(operand.unwrap())
- self.init(<shared_ptr[CExpression]> expression)
+ cdef shared_ptr[CNotExpression] expr
+ expr = CMakeNotExpression(operand.unwrap())
+ self.init(<shared_ptr[CExpression]> expr)
cdef class AndExpression(BinaryExpression):
- def __init__(self, Expression left_operand not None,
- Expression right_operand not None,
+ def __init__(self, Expression left not None, Expression right not None,
*additional_operands):
cdef:
Expression operand
vector[shared_ptr[CExpression]] exprs
- exprs.push_back(left_operand.unwrap())
- exprs.push_back(right_operand.unwrap())
+ exprs.push_back(left.unwrap())
+ exprs.push_back(right.unwrap())
for operand in additional_operands:
exprs.push_back(operand.unwrap())
- self.init(MakeAndExpression(exprs))
+ self.init(CMakeAndExpression(exprs))
cdef class OrExpression(BinaryExpression):
- def __init__(self, Expression left_operand not None,
- Expression right_operand not None,
+ def __init__(self, Expression left not None, Expression right not None,
*additional_operands):
cdef:
Expression operand
vector[shared_ptr[CExpression]] exprs
- exprs.push_back(left_operand.unwrap())
- exprs.push_back(right_operand.unwrap())
+ exprs.push_back(left.unwrap())
+ exprs.push_back(right.unwrap())
for operand in additional_operands:
exprs.push_back(operand.unwrap())
- self.init(MakeOrExpression(exprs))
+ self.init(CMakeOrExpression(exprs))
diff --git a/python/pyarrow/dataset.py b/python/pyarrow/dataset.py
index 2d8a458..1aa0dfb 100644
--- a/python/pyarrow/dataset.py
+++ b/python/pyarrow/dataset.py
@@ -34,7 +34,6 @@ from pyarrow._dataset import ( # noqa
ComparisonExpression,
Dataset,
DatasetFactory,
- DefaultPartitioning,
DirectoryPartitioning,
Expression,
FieldExpression,
@@ -52,7 +51,6 @@ from pyarrow._dataset import ( # noqa
PartitioningFactory,
ScalarExpression,
Scanner,
- ScannerBuilder,
ScanTask,
Source,
TreeSource,
@@ -332,5 +330,38 @@ def dataset(sources, filesystem=None, partitioning=None,
format=None):
format=format)
for src in sources
]
-
return DatasetFactory(sources).finish()
+
+
+def field(name):
+ """References a named column of the dataset.
+
+ Stores only the field's name. Type and other information is known only when
+ the expression is applied on a dataset having an explicit scheme.
+
+ Parameters
+ ----------
+ name : string
+ The name of the field the expression references to.
+
+ Returns
+ -------
+ field_expr : FieldExpression
+ """
+ return FieldExpression(name)
+
+
+def scalar(value):
+ """Expression representing a scalar value.
+
+ Parameters
+ ----------
+ value : bool, int, float or string
+ Python value of the scalar. Note that only a subset of types are
+ currently supported.
+
+ Returns
+ -------
+ scalar_expr : ScalarExpression
+ """
+ return ScalarExpression(value)
diff --git a/python/pyarrow/includes/libarrow.pxd
b/python/pyarrow/includes/libarrow.pxd
index d4cd5f2..7212896 100644
--- a/python/pyarrow/includes/libarrow.pxd
+++ b/python/pyarrow/includes/libarrow.pxd
@@ -800,7 +800,11 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
cdef cppclass CDoubleScalar" arrow::DoubleScalar"(CScalar):
double value
+ cdef cppclass CStringScalar" arrow::StringScalar"(CScalar):
+ pass
+
shared_ptr[CScalar] MakeScalar[Value](Value value)
+ shared_ptr[CScalar] MakeStringScalar" arrow::MakeScalar"(c_string value)
cdef cppclass CConcatenateTablesOptions" arrow::ConcatenateTablesOptions":
c_bool unify_schemas
diff --git a/python/pyarrow/includes/libarrow_dataset.pxd
b/python/pyarrow/includes/libarrow_dataset.pxd
index 0c07aaa..10e3793 100644
--- a/python/pyarrow/includes/libarrow_dataset.pxd
+++ b/python/pyarrow/includes/libarrow_dataset.pxd
@@ -120,11 +120,11 @@ cdef extern from "arrow/dataset/api.h" namespace
"arrow::dataset" nogil:
CUnaryExpression):
CInExpression(shared_ptr[CExpression] operand, shared_ptr[CArray] set)
- cdef shared_ptr[CNotExpression] MakeNotExpression "arrow::dataset::not_"(
+ cdef shared_ptr[CNotExpression] CMakeNotExpression "arrow::dataset::not_"(
shared_ptr[CExpression] operand)
- cdef shared_ptr[CExpression] MakeAndExpression "arrow::dataset::and_"(
+ cdef shared_ptr[CExpression] CMakeAndExpression "arrow::dataset::and_"(
const CExpressionVector& subexpressions)
- cdef shared_ptr[CExpression] MakeOrExpression "arrow::dataset::or_"(
+ cdef shared_ptr[CExpression] CMakeOrExpression "arrow::dataset::or_"(
const CExpressionVector& subexpressions)
cdef CResult[shared_ptr[CExpression]] CInsertImplicitCasts \
@@ -141,8 +141,6 @@ cdef extern from "arrow/dataset/api.h" namespace
"arrow::dataset" nogil:
shared_ptr[CExpression] filter
shared_ptr[CSchema] schema
c_bool use_threads
- # shared_ptr[CExpressionEvaluator] evaluator
- # shared_ptr[CRecordBatchProjector] projector
@staticmethod
shared_ptr[CScanOptions] Defaults()
@@ -153,10 +151,6 @@ cdef extern from "arrow/dataset/api.h" namespace
"arrow::dataset" nogil:
cdef cppclass CScanTask" arrow::dataset::ScanTask":
CResult[CRecordBatchIterator] Execute()
- cdef cppclass CInMemoryScanTask "arrow::dataset::InMemoryScanTask"(
- CScanTask):
- pass
-
ctypedef CIterator[shared_ptr[CScanTask]] CScanTaskIterator \
"arrow::dataset::ScanTaskIterator"
@@ -311,10 +305,6 @@ cdef extern from "arrow/dataset/api.h" namespace
"arrow::dataset" nogil:
cdef cppclass CPartitioningFactory "arrow::dataset::PartitioningFactory":
pass
- cdef cppclass CDefaultPartitioning \
- "arrow::dataset::DefaultPartitioning"(CPartitioning):
- CDefaultPartitioning()
-
cdef cppclass CDirectoryPartitioning \
"arrow::dataset::DirectoryPartitioning"(CPartitioning):
CDirectoryPartitioning(shared_ptr[CSchema] schema)
diff --git a/python/pyarrow/tests/test_dataset.py
b/python/pyarrow/tests/test_dataset.py
index fe3eb22..d5f40cb 100644
--- a/python/pyarrow/tests/test_dataset.py
+++ b/python/pyarrow/tests/test_dataset.py
@@ -16,9 +16,9 @@
# under the License.
import sys
+import operator
import numpy as np
-
import pytest
import pyarrow as pa
@@ -175,13 +175,13 @@ def test_filesystem_source(mockfs):
partitions = [ds.ScalarExpression(True), ds.ScalarExpression(True)]
source = ds.FileSystemSource(schema,
- source_partition=None,
+ root_partition=None,
file_format=file_format,
filesystem=mockfs,
paths_or_selector=paths,
partitions=partitions)
- source_partition = ds.ComparisonExpression(
+ root_partition = ds.ComparisonExpression(
ds.CompareOperator.Equal,
ds.FieldExpression('source'),
ds.ScalarExpression(1337)
@@ -199,42 +199,35 @@ def test_filesystem_source(mockfs):
)
]
source = ds.FileSystemSource(paths_or_selector=paths, schema=schema,
- source_partition=source_partition,
+ root_partition=root_partition,
filesystem=mockfs, partitions=partitions,
file_format=file_format)
- assert source.partition_expression.equals(source_partition)
+ assert source.partition_expression.equals(root_partition)
def test_dataset(dataset):
assert isinstance(dataset, ds.Dataset)
assert isinstance(dataset.schema, pa.Schema)
- # TODO(kszucs): test non-boolean expressions for filter do raise
- builder = dataset.new_scan()
- assert isinstance(builder, ds.ScannerBuilder)
-
- scanner = builder.finish()
- assert isinstance(scanner, ds.Scanner)
- assert len(list(scanner.scan())) == 2
+ # TODO(kszucs): test non-boolean Exprs for filter do raise
expected_i64 = pa.array([0, 1, 2, 3, 4], type=pa.int64())
expected_f64 = pa.array([0, 1, 2, 3, 4], type=pa.float64())
- for task in scanner.scan():
+ for task in dataset.scan():
assert isinstance(task, ds.ScanTask)
for batch in task.execute():
assert batch.column(0).equals(expected_i64)
assert batch.column(1).equals(expected_f64)
- table = scanner.to_table()
+ batches = dataset.to_batches()
+ assert all(isinstance(batch, pa.RecordBatch) for batch in batches)
+
+ table = dataset.to_table()
assert isinstance(table, pa.Table)
assert len(table) == 10
- condition = ds.ComparisonExpression(
- ds.CompareOperator.Equal,
- ds.FieldExpression('i64'),
- ds.ScalarExpression(1)
- )
- scanner = dataset.new_scan().use_threads(True).filter(condition).finish()
+ condition = ds.field('i64') == 1
+ scanner = ds.Scanner(dataset, use_threads=True, filter=condition)
result = scanner.to_table().to_pydict()
# don't rely on the scanning order
@@ -244,17 +237,16 @@ def test_dataset(dataset):
assert sorted(result['key']) == ['xxx', 'yyy']
-def test_scanner_builder(dataset):
- builder = ds.ScannerBuilder(dataset, memory_pool=pa.default_memory_pool())
- scanner = builder.finish()
+def test_scanner(dataset):
+ scanner = ds.Scanner(dataset, memory_pool=pa.default_memory_pool())
assert isinstance(scanner, ds.Scanner)
assert len(list(scanner.scan())) == 2
with pytest.raises(pa.ArrowInvalid):
- dataset.new_scan().project(['unknown'])
+ dataset.scan(columns=['unknown'])
- builder = dataset.new_scan(memory_pool=pa.default_memory_pool())
- scanner = builder.project(['i64']).finish()
+ scanner = ds.Scanner(dataset, columns=['i64'],
+ memory_pool=pa.default_memory_pool())
assert isinstance(scanner, ds.Scanner)
assert len(list(scanner.scan())) == 2
@@ -294,18 +286,7 @@ def test_partitioning():
expr = partitioning.parse('/3/3.14')
assert isinstance(expr, ds.Expression)
- expected = ds.AndExpression(
- ds.ComparisonExpression(
- ds.CompareOperator.Equal,
- ds.FieldExpression('group'),
- ds.ScalarExpression(3)
- ),
- ds.ComparisonExpression(
- ds.CompareOperator.Equal,
- ds.FieldExpression('key'),
- ds.ScalarExpression(3.14)
- )
- )
+ expected = (ds.field('group') == 3) & (ds.field('key') == 3.14)
assert expr.equals(expected)
with pytest.raises(pa.ArrowInvalid):
@@ -318,17 +299,9 @@ def test_partitioning():
])
)
expr = partitioning.parse('/alpha=0/beta=3')
- expected = ds.AndExpression(
- ds.ComparisonExpression(
- ds.CompareOperator.Equal,
- ds.FieldExpression('alpha'),
- ds.ScalarExpression(0)
- ),
- ds.ComparisonExpression(
- ds.CompareOperator.Equal,
- ds.FieldExpression('beta'),
- ds.ScalarExpression(3)
- )
+ expected = (
+ (ds.field('alpha') == ds.scalar(0)) &
+ (ds.field('beta') == ds.scalar(3))
)
assert expr.equals(expected)
@@ -337,6 +310,7 @@ def test_expression():
a = ds.ScalarExpression(1)
b = ds.ScalarExpression(1.1)
c = ds.ScalarExpression(True)
+ d = ds.ScalarExpression("string")
equal = ds.ComparisonExpression(ds.CompareOperator.Equal, a, b)
assert equal.op() == ds.CompareOperator.Equal
@@ -349,7 +323,7 @@ def test_expression():
ds.AndExpression(a, b, c)
ds.OrExpression(a, b)
- ds.OrExpression(a, b, c)
+ ds.OrExpression(a, b, c, d)
ds.NotExpression(ds.OrExpression(a, b, c))
ds.IsValidExpression(a)
ds.CastExpression(a, pa.int32())
@@ -382,6 +356,67 @@ def test_expression():
assert str(condition) == "(i64 > 5:int64)"
+def test_expression_ergonomics():
+ zero = ds.scalar(0)
+ one = ds.scalar(1)
+ true = ds.scalar(True)
+ false = ds.scalar(False)
+ string = ds.scalar("string")
+ field = ds.field("field")
+
+ assert one.equals(ds.ScalarExpression(1))
+ assert zero.equals(ds.ScalarExpression(0))
+ assert true.equals(ds.ScalarExpression(True))
+ assert false.equals(ds.ScalarExpression(False))
+ assert string.equals(ds.ScalarExpression("string"))
+ assert field.equals(ds.FieldExpression("field"))
+
+ expected = ds.AndExpression(ds.ScalarExpression(1), ds.ScalarExpression(0))
+ for expr in [one & zero, 1 & zero, one & 0]:
+ assert expr.equals(expected)
+
+ expected = ds.OrExpression(ds.ScalarExpression(1), ds.ScalarExpression(0))
+ for expr in [one | zero, 1 | zero, one | 0]:
+ assert expr.equals(expected)
+
+ comparison_ops = [
+ (operator.eq, ds.CompareOperator.Equal),
+ (operator.ne, ds.CompareOperator.NotEqual),
+ (operator.ge, ds.CompareOperator.GreaterEqual),
+ (operator.le, ds.CompareOperator.LessEqual),
+ (operator.lt, ds.CompareOperator.Less),
+ (operator.gt, ds.CompareOperator.Greater),
+ ]
+ for op, compare_op in comparison_ops:
+ expr = op(zero, one)
+ expected = ds.ComparisonExpression(compare_op, zero, one)
+ assert expr.equals(expected)
+
+ expr = ~true == false
+ expected = ds.ComparisonExpression(
+ ds.CompareOperator.Equal,
+ ds.NotExpression(ds.ScalarExpression(True)),
+ ds.ScalarExpression(False)
+ )
+ assert expr.equals(expected)
+
+ for typ in ("bool", pa.bool_()):
+ expr = field.cast(typ) == true
+ expected = ds.ComparisonExpression(
+ ds.CompareOperator.Equal,
+ ds.CastExpression(ds.FieldExpression("field"), pa.bool_()),
+ ds.ScalarExpression(True)
+ )
+ assert expr.equals(expected)
+
+ expr = field.isin([1, 2])
+ expected = ds.InExpression(ds.FieldExpression("field"), pa.array([1, 2]))
+ assert expr.equals(expected)
+
+ with pytest.raises(TypeError):
+ field.isin(1)
+
+
@pytest.mark.parametrize('paths_or_selector', [
fs.FileSelector('subdir', recursive=True),
[
@@ -423,10 +458,9 @@ def test_file_system_factory(mockfs, paths_or_selector):
assert isinstance(source, ds.Source)
dataset = ds.Dataset([source], inspected_schema)
+ assert len(list(dataset.scan())) == 2
- scanner = dataset.new_scan().finish()
- assert len(list(scanner.scan())) == 2
-
+ scanner = ds.Scanner(dataset)
expected_i64 = pa.array([0, 1, 2, 3, 4], type=pa.int64())
expected_f64 = pa.array([0, 1, 2, 3, 4], type=pa.float64())
for task, group, key in zip(scanner.scan(), [1, 2], ['xxx', 'yyy']):
@@ -439,7 +473,7 @@ def test_file_system_factory(mockfs, paths_or_selector):
assert batch[2].equals(expected_group_column)
assert batch[3].equals(expected_key_column)
- table = scanner.to_table()
+ table = dataset.to_table()
assert isinstance(table, pa.Table)
assert len(table) == 10
assert table.num_columns == 4
@@ -533,19 +567,19 @@ def _check_dataset_from_path(path, table, **kwargs):
assert isinstance(path, pathlib.Path)
dataset = ds.dataset(ds.source(path, **kwargs))
assert dataset.schema.equals(table.schema, check_metadata=False)
- result = dataset.new_scan().finish().to_table()
+ result = dataset.to_table()
assert result.replace_schema_metadata().equals(table)
# string path
dataset = ds.dataset(ds.source(str(path), **kwargs))
assert dataset.schema.equals(table.schema, check_metadata=False)
- result = dataset.new_scan().finish().to_table()
+ result = dataset.to_table()
assert result.replace_schema_metadata().equals(table)
# passing directly to dataset
dataset = ds.dataset(str(path), **kwargs)
assert dataset.schema.equals(table.schema, check_metadata=False)
- result = dataset.new_scan().finish().to_table()
+ result = dataset.to_table()
assert result.replace_schema_metadata().equals(table)
@@ -573,7 +607,7 @@ def test_open_dataset_list_of_files(tempdir):
ds.dataset(ds.source([path1, path2])),
ds.dataset(ds.source([str(path1), str(path2)]))]:
assert dataset.schema.equals(table.schema, check_metadata=False)
- result = dataset.new_scan().finish().to_table()
+ result = dataset.to_table()
assert result.replace_schema_metadata().equals(table)
@@ -609,7 +643,7 @@ def test_open_dataset_partitioned_directory(tempdir):
expected_schema = table.schema.append(pa.field("part", pa.int8()))
assert dataset.schema.equals(expected_schema, check_metadata=False)
- result = dataset.new_scan().finish().to_table()
+ result = dataset.to_table()
expected = full_table.append_column(
"part", pa.array(np.repeat([0, 1, 2], 9), type=pa.int8()))
assert result.replace_schema_metadata().equals(expected)
@@ -662,15 +696,9 @@ def test_filter_implicit_cast(tempdir):
_, path = _create_single_file(tempdir, table)
dataset = ds.dataset(str(path))
- filter_ = ds.ComparisonExpression(
- ds.CompareOperator.Greater,
- ds.FieldExpression('a'),
- ds.ScalarExpression(2)
- )
-
- scanner_builder = dataset.new_scan()
- scanner_builder.filter(filter_)
- result = scanner_builder.finish().to_table()
+ filter_ = ds.field('a') > 2
+ scanner = ds.Scanner(dataset, filter=filter_)
+ result = scanner.to_table()
assert len(result) == 3
