Expose more methods

cpp/src/arrow/dataset/dataset.h

@@ -151,7 +151,7 @@ class ARROW_DS_EXPORT Dataset : public std::enable_shared_from_this<Dataset> {
   /// \brief Build a Dataset from uniform sources.
   //
   /// \param[in] sources one or more input data sources
-  /// \param[in] schema a known schema to conform to, may be nullptr
+  /// \param[in] schema a known schema to conform to
   static Result<DatasetPtr> Make(DataSourceVector sources,
                                  std::shared_ptr<Schema> schema);
 

python/pyarrow/_dataset.pyx

@@ -24,27 +24,18 @@ import six
 
 from cython.operator cimport dereference as deref
 
-import pyarrow as pa
-from pyarrow.compat import frombytes, tobytes
 from pyarrow.lib cimport *
-# from pyarrow.lib import ArrowException
-# from pyarrow.lib import as_buffer
-# from pyarrow.lib cimport RecordBatch, MemoryPool, Schema, check_status
-from pyarrow._fs cimport FileSystem, Selector
 from pyarrow.includes.libarrow_dataset cimport *
+from pyarrow.compat import frombytes, tobytes
+from pyarrow._fs cimport FileSystem, Selector
 
 
-__all__ = [
-    'ScanOptions',
-    'ScanContext',
-    'ScanTask',
-    'SimpleScanTask',
-    'DataFragment',
-    'SimpleDataFragment',
-    'DataSource',
-    'SimpleDataSource',
-    'Dataset'
-]
+def _forbid_instantiation(klass):
+    subclasses = [cls.__name__ for cls in klass.__subclasses__]
+    raise TypeError(
+        '{} is an abstract class thus cannot be initialized. Use one of the '
+        'subclasses instead: {}'.format(klass.__name__, ', '.join(subclasses))
+    )
 
 
 cdef class WriteOptions:
@@ -73,6 +64,9 @@ cdef class FileFormat:
         shared_ptr[CFileFormat] wrapped
         CFileFormat* format
 
+    def __init__(self):
+        _forbid_instantiation(self.__class__)
+
     cdef void init(self, const shared_ptr[CFileFormat]& sp):
         self.wrapped = sp
         self.format = sp.get()
@@ -204,6 +198,9 @@ cdef class DataSourceDiscovery:
         shared_ptr[CDataSourceDiscovery] wrapped
         CDataSourceDiscovery* discovery
 
+    def __init__(self):
+        _forbid_instantiation(self.__class__)
+
     cdef init(self, shared_ptr[CDataSourceDiscovery]& sp):
         self.wrapped = sp
         self.discovery = sp.get()
@@ -310,7 +307,7 @@ cdef class FileSource:
         try:
             return self.equals(other)
         except TypeError:
-            return NotImplementedError
+            return NotImplemented
 
     @property
     def compression(self):
@@ -329,6 +326,9 @@ cdef class DataFragment:
         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()
@@ -404,6 +404,9 @@ cdef class DataSource:
         shared_ptr[CDataSource] wrapped
         CDataSource* source
 
+    def __init__(self):
+        _forbid_instantiation(self.__class__)
+
     cdef void init(self, const shared_ptr[CDataSource]& sp):
         self.wrapped = sp
         self.source = sp.get()
@@ -428,10 +431,6 @@ cdef class DataSource:
     cdef shared_ptr[CDataSource] unwrap(self):
         return self.wrapped
 
-    # def __iter__(self):
-    #     for fragment in self.fragments():
-    #         yield fragment
-
     def fragments(self, ScanOptions options=None):
         cdef:
             CDataFragmentIterator iterator
@@ -524,20 +523,18 @@ cdef class Dataset:
         shared_ptr[CDataset] wrapped
         CDataset* dataset
 
-    def __init__(self, data_sources, Schema schema=None):
+    def __init__(self, data_sources, Schema schema not None):
         cdef:
             DataSource source
             CDataSourceVector sources
             CResult[CDatasetPtr] result
+            shared_ptr[CSchema] sp_schema
 
         for source in data_sources:
             sources.push_back(source.wrapped)
 
-        result = CDataset.Make(sources, schema.sp_schema)
-        # if schema is None:
-        #     dataset = make_shared[CDataset](sources)
-        # else:
-        #     dataset = make_shared[CDataset](sources, schema.sp_schema)
+        result = CDataset.Make(sources, pyarrow_unwrap_schema(schema))
+
         self.init(GetResultValue(result))
 
     cdef void init(self, const shared_ptr[CDataset]& sp):
@@ -551,7 +548,6 @@ cdef class Dataset:
     def new_scan(self):
         cdef shared_ptr[CScannerBuilder] builder
         builder = GetResultValue(self.dataset.NewScan())
-        # return ScannerBuilder(self, context)
         return ScannerBuilder.wrap(builder)
 
     @property
@@ -700,6 +696,13 @@ cdef class ScannerBuilder:
         check_status(self.builder.Filter(expression.unwrap()))
         return self
 
+    def use_threads(self, bint value):
+        check_status(self.builder.UseThreads(value))
+        return self
+
+    def schema(self):
+        return pyarrow_wrap_schema(self.builder.schema())
+
 
 cdef class Scanner:
 
@@ -761,14 +764,14 @@ cdef class Expression:
         CExpression* expression
 
     def __init__(self):
-        raise TypeError('Do not ititialize')
+        _forbid_instantiation(self.__class__)
 
     cdef void init(self, const shared_ptr[CExpression]& sp):
         self.wrapped = sp
         self.expression = sp.get()
 
     @staticmethod
-    cdef wrap(shared_ptr[CExpression]& sp):
+    cdef wrap(const shared_ptr[CExpression]& sp):
         cdef Expression self
 
         typ = sp.get().type()
@@ -801,6 +804,29 @@ cdef class Expression:
     cdef inline shared_ptr[CExpression] unwrap(self):
         return self.wrapped
 
+    def equals(self, Expression other):
+        return self.expression.Equals(other.unwrap())
+
+    def __str__(self):
+        return frombytes(self.expression.ToString())
+
+    def __eq__(self, other):
+        try:
+            return self.equals(other)
+        except TypeError:
+            return NotImplemented
+
+    def validate(self, Schema schema not None):
+        cdef:
+            shared_ptr[CSchema] sp_schema
+            CResult[shared_ptr[CDataType]] result
+        sp_schema = pyarrow_unwrap_schema(schema)
+        result = self.expression.Validate(deref(sp_schema))
+        return pyarrow_wrap_data_type(GetResultValue(result))
+
+    def assume(self, Expression given):
+        return Expression.wrap(self.expression.Assume(given.unwrap()))
+
 
 cdef class UnaryExpression(Expression):
 
@@ -819,6 +845,14 @@ cdef class BinaryExpression(Expression):
         Expression.init(self, sp)
         self.binary = <CBinaryExpression*> sp.get()
 
+    @property
+    def left_operand(self):
+        return Expression.wrap(self.binary.left_operand())
+
+    @property
+    def right_operand(self):
+        return Expression.wrap(self.binary.right_operand())
+
 
 cdef class ScalarExpression(Expression):
 
@@ -845,6 +879,7 @@ cdef class ScalarExpression(Expression):
         Expression.init(self, sp)
         self.scalar = <CScalarExpression*> sp.get()
 
+    # @property
     # def value(self):
     #     return pyarrow_wrap_scalar(self.scalar.value())
 
@@ -853,6 +888,13 @@ cdef class FieldExpression(Expression):
 
     cdef CFieldExpression* scalar
 
+    def __init__(self, name):
+        cdef:
+            c_string field_name = tobytes(name)
+            shared_ptr[CExpression] expression
+        expression.reset(new CFieldExpression(field_name))
+        self.init(expression)
+
     cdef void init(self, const shared_ptr[CExpression]& sp):
         Expression.init(self, sp)
         self.scalar = <CFieldExpression*> sp.get()
@@ -891,8 +933,8 @@ cdef class ComparisonExpression(BinaryExpression):
         BinaryExpression.init(self, sp)
         self.comparison = <CComparisonExpression*> sp.get()
 
-    # def op(self):
-    #     return ...
+    def op(self):
+        return <CompareOperator> self.comparison.op()
 
 
 cdef class NotExpression(UnaryExpression):

python/pyarrow/dataset.py

@@ -22,12 +22,14 @@ import sys
 if sys.version_info < (3,):
     raise ImportError("Python Dataset bindings require Python 3")
 
-from pyarrow._dataset import (
+from pyarrow._dataset import (  # noqa
     Dataset,
     DataSource,
     DataFragment,
     ParquetFileFormat,
     FileSource,
+    FileFormat,
+    Expression,
     SimpleDataFragment,
     ScanOptions,
     ScanTask,

python/pyarrow/includes/libarrow_dataset.pxd

@@ -47,8 +47,10 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CExpressionType_CAST "arrow::dataset::ExpressionType::type::CAST"
         CExpressionType_AND "arrow::dataset::ExpressionType::type::AND"
         CExpressionType_OR "arrow::dataset::ExpressionType::type::OR"
-        CExpressionType_COMPARISON "arrow::dataset::ExpressionType::type::COMPARISON"
-        CExpressionType_IS_VALID "arrow::dataset::ExpressionType::type::IS_VALID"
+        CExpressionType_COMPARISON \
+            "arrow::dataset::ExpressionType::type::COMPARISON"
+        CExpressionType_IS_VALID \
+            "arrow::dataset::ExpressionType::type::IS_VALID"
         CExpressionType_IN "arrow::dataset::ExpressionType::type::IN"
         CExpressionType_CUSTOM "arrow::dataset::ExpressionType::type::CUSTOM"
 
@@ -58,61 +60,80 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         c_bool Equals(const shared_ptr[CExpression]& other) const
         c_bool IsNull() const
         CResult[shared_ptr[CDataType]] Validate(const CSchema& schema) const
-
         shared_ptr[CExpression] Assume(const CExpression& given) const
-        shared_ptr[CExpression] Assume(const shared_ptr[CExpression]& given) const
+        shared_ptr[CExpression] Assume(
+            const shared_ptr[CExpression]& given) const
         c_string ToString() const
         CExpressionType type() const
         shared_ptr[CExpression] Copy() const
+        # TODO(kszucs): expose the following methods
         # InExpression In(shared_ptr[CArray] set) const;
         # IsValidExpression IsValid() const
-        # CastExpression CastTo(shared_ptr[CDataType] type,
-        #                       compute::CastOptions options = compute::CastOptions()) const;
-        # CastExpression CastLike(const CExpression& expr,
-        #                         compute::CastOptions options = compute::CastOptions()) const;
-        # CastExpression CastLike(shared_ptr[CExpression] expr,
-        #                         compute::CastOptions options = compute::CastOptions()) const;
+        # CastExpression CastTo(
+        #     shared_ptr[CDataType] type,
+        #     compute::CastOptions options = compute::CastOptions()
+        # ) const;
+        # CastExpression CastLike(
+        #     const CExpression& expr,
+        #     compute::CastOptions options = compute::CastOptions()
+        # ) const;
+        # CastExpression CastLike(
+        #     shared_ptr[CExpression] expr,
+        #     compute::CastOptions options = compute::CastOptions()
+        # ) const;
 
-    ctypedef shared_ptr[CExpression] CExpressionPtr "arrow::dataset::ExpressionPtr"
-    ctypedef vector[CExpressionPtr] CExpressionVector "arrow::dataset::ExpressionVector"
+    ctypedef shared_ptr[CExpression] CExpressionPtr \
+        "arrow::dataset::ExpressionPtr"
+    ctypedef vector[CExpressionPtr] CExpressionVector \
+        "arrow::dataset::ExpressionVector"
 
-    cdef cppclass CUnaryExpression "arrow::dataset::UnaryExpression"(CExpression):
+    cdef cppclass CUnaryExpression "arrow::dataset::UnaryExpression"(
+            CExpression):
         const CExpressionPtr& operand() const
 
-    cdef cppclass CBinaryExpression "arrow::dataset::BinaryExpression"(CExpression):
+    cdef cppclass CBinaryExpression "arrow::dataset::BinaryExpression"(
+            CExpression):
         const CExpressionPtr& left_operand() const
         const CExpressionPtr& right_operand() const
 
-    cdef cppclass CScalarExpression "arrow::dataset::ScalarExpression"(CExpression):
+    cdef cppclass CScalarExpression "arrow::dataset::ScalarExpression"(
+            CExpression):
         CScalarExpression(const shared_ptr[CScalar]& value)
         const shared_ptr[CScalar]& value() const
 
-    cdef cppclass CFieldExpression "arrow::dataset::FieldExpression"(CExpression):
+    cdef cppclass CFieldExpression "arrow::dataset::FieldExpression"(
+            CExpression):
+        CFieldExpression(c_string name)
         c_string name() const
 
-    cdef cppclass CComparisonExpression "arrow::dataset::ComparisonExpression"(CBinaryExpression):
+    cdef cppclass CComparisonExpression "arrow::dataset::ComparisonExpression"(
+            CBinaryExpression):
         CComparisonExpression(CCompareOperator op, CExpressionPtr left_operand,
                               CExpressionPtr right_operand)
         CCompareOperator op() const
 
-    cdef cppclass CAndExpression "arrow::dataset::AndExpression"(CBinaryExpression):
+    cdef cppclass CAndExpression "arrow::dataset::AndExpression"(
+            CBinaryExpression):
         pass
 
-    cdef cppclass COrExpression "arrow::dataset::OrExpression"(CBinaryExpression):
+    cdef cppclass COrExpression "arrow::dataset::OrExpression"(
+            CBinaryExpression):
         pass
 
-    cdef cppclass CNotExpression "arrow::dataset::NotExpression"(CBinaryExpression):
+    cdef cppclass CNotExpression "arrow::dataset::NotExpression"(
+            CBinaryExpression):
         pass
 
-    cdef cppclass CIsValidExpression "arrow::dataset::IsValidExpression"(CBinaryExpression):
+    cdef cppclass CIsValidExpression "arrow::dataset::IsValidExpression"(
+            CBinaryExpression):
         pass
 
-    cdef shared_ptr[CNotExpression] MakeNotExpression "arrow::dataset::not_"(CExpressionPtr operand)
-    cdef CExpressionPtr MakeAndExpression "arrow::dataset::and_"(const CExpressionVector& subexpressions)
-    cdef CExpressionPtr MakeOrExpression "arrow::dataset::or_"(const CExpressionVector& subexpressions)
-
-    # InExpression
-    # CastExpression
+    cdef shared_ptr[CNotExpression] MakeNotExpression "arrow::dataset::not_"(
+            CExpressionPtr operand)
+    cdef CExpressionPtr MakeAndExpression "arrow::dataset::and_"(
+            const CExpressionVector& subexpressions)
+    cdef CExpressionPtr MakeOrExpression "arrow::dataset::or_"(
+            const CExpressionVector& subexpressions)
 
     cdef cppclass CFilter "arrow::dataset::Filter":
         pass
@@ -134,8 +155,10 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     cdef cppclass CScanContext "arrow::dataset::ScanContext":
         CMemoryPool* pool
 
-    ctypedef shared_ptr[CScanOptions] CScanOptionsPtr "arrow::dataset::ScanOptionsPtr"
-    ctypedef shared_ptr[CScanContext] CScanContextPtr "arrow::dataset::ScanContextPtr"
+    ctypedef shared_ptr[CScanOptions] CScanOptionsPtr \
+        "arrow::dataset::ScanOptionsPtr"
+    ctypedef shared_ptr[CScanContext] CScanContextPtr \
+        "arrow::dataset::ScanContextPtr"
 
     cdef cppclass CScanTask" arrow::dataset::ScanTask":
         CResult[CRecordBatchIterator] Scan()
@@ -144,7 +167,8 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         pass
 
     ctypedef shared_ptr[CScanTask] CScanTaskPtr "arrow::dataset::ScanTaskPtr"
-    ctypedef CIterator[CScanTaskPtr] CScanTaskIterator "arrow::dataset::ScanTaskIterator"
+    ctypedef CIterator[CScanTaskPtr] CScanTaskIterator \
+        "arrow::dataset::ScanTaskIterator"
 
     cdef cppclass CScanner "arrow::dataset::Scanner":
         CResult[CScanTaskIterator] Scan()
@@ -161,16 +185,20 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CResult[CScannerPtr] Finish()
         shared_ptr[CSchema] schema() const
 
-    ctypedef shared_ptr[CScannerBuilder] CScannerBuilderPtr "arrow::dataset::ScannerBuilderPtr"
+    ctypedef shared_ptr[CScannerBuilder] CScannerBuilderPtr \
+        "arrow::dataset::ScannerBuilderPtr"
 
     cdef cppclass CDataFragment "arrow::dataset::DataFragment":
         CResult[CScanTaskIterator] Scan(CScanContextPtr context)
         c_bool splittable()
         CScanOptionsPtr scan_options()
 
-    ctypedef shared_ptr[CDataFragment] CDataFragmentPtr "arrow::dataset::DataFragmentPtr"
-    ctypedef vector[CDataFragmentPtr] CDataFragmentVector "arrow::dataset::DataFragmentVector"
-    ctypedef CIterator[CDataFragmentPtr] CDataFragmentIterator "arrow::dataset::DataFragmentIterator"
+    ctypedef shared_ptr[CDataFragment] CDataFragmentPtr \
+        "arrow::dataset::DataFragmentPtr"
+    ctypedef vector[CDataFragmentPtr] CDataFragmentVector \
+        "arrow::dataset::DataFragmentVector"
+    ctypedef CIterator[CDataFragmentPtr] CDataFragmentIterator \
+        "arrow::dataset::DataFragmentIterator"
 
     cdef cppclass CSimpleDataFragment "arrow::dataset::SimpleDataFragment"(
             CDataFragment):
@@ -182,14 +210,18 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         const CExpressionPtr& partition_expression()
         c_string type()
 
-    cdef cppclass CSimpleDataSource "arrow::dataset::SimpleDataSource"(CDataSource):
+    cdef cppclass CSimpleDataSource "arrow::dataset::SimpleDataSource"(
+            CDataSource):
         pass
 
-    cdef cppclass CTreeDataSource "arrow::dataset::TreeDataSource"(CDataSource):
+    cdef cppclass CTreeDataSource "arrow::dataset::TreeDataSource"(
+            CDataSource):
         pass
 
-    ctypedef shared_ptr[CDataSource] CDataSourcePtr "arrow::dataset::DataSourcePtr"
-    ctypedef vector[CDataSourcePtr] CDataSourceVector "arrow::dataset::DataSourceVector"
+    ctypedef shared_ptr[CDataSource] CDataSourcePtr \
+        "arrow::dataset::DataSourcePtr"
+    ctypedef vector[CDataSourcePtr] CDataSourceVector \
+        "arrow::dataset::DataSourceVector"
 
     cdef cppclass CDataset "arrow::dataset::Dataset":
         @staticmethod
@@ -235,7 +267,8 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
                              CScanOptionsPtr opts,
                              CDataFragmentPtr* out)
 
-    ctypedef shared_ptr[CFileFormat] CFileFormatPtr "arrow::dataset::FileFormatPtr"
+    ctypedef shared_ptr[CFileFormat] CFileFormatPtr \
+        "arrow::dataset::FileFormatPtr"
 
     cdef cppclass CFileBasedDataFragment \
             "arrow::dataset::FileBasedDataFragment"(CDataFragment):
@@ -285,19 +318,23 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         c_string name() const
         CResult[CExpressionPtr] Parse(const c_string& path) const
 
-    ctypedef shared_ptr[CPartitionScheme] CPartitionSchemePtr "arrow::dataset::PartitionSchemePtr"
+    ctypedef shared_ptr[CPartitionScheme] CPartitionSchemePtr \
+        "arrow::dataset::PartitionSchemePtr"
 
-    cdef cppclass CSchemaPartitionScheme "arrow::dataset::SchemaPartitionScheme"(CPartitionScheme):
+    cdef cppclass CSchemaPartitionScheme \
+            "arrow::dataset::SchemaPartitionScheme"(CPartitionScheme):
         CSchemaPartitionScheme(shared_ptr[CSchema] schema)
         const shared_ptr[CSchema]& schema()
 
-    cdef cppclass CHivePartitionScheme "arrow::dataset::HivePartitionScheme"(CPartitionScheme):
+    cdef cppclass CHivePartitionScheme \
+            "arrow::dataset::HivePartitionScheme"(CPartitionScheme):
         CHivePartitionScheme(shared_ptr[CSchema] schema)
         # vector[CUnconvertedKey] GetUnconvertedKeys(const c_string& path) const;
 
     ############################### Discovery #################################
 
-    cdef cppclass CFileSystemDiscoveryOptions "arrow::dataset::FileSystemDiscoveryOptions":
+    cdef cppclass CFileSystemDiscoveryOptions \
+            "arrow::dataset::FileSystemDiscoveryOptions":
         c_string partition_base_dir
         c_bool exclude_invalid_files
         vector[c_string] ignore_prefixes
@@ -312,16 +349,23 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CExpressionPtr root_partition()
         CStatus SetRootPartition(CExpressionPtr partition)
 
-    ctypedef shared_ptr[CDataSourceDiscovery] CDataSourceDiscoveryPtr "arrow::dataset::DataSourceDiscovery"
+    ctypedef shared_ptr[CDataSourceDiscovery] CDataSourceDiscoveryPtr \
+        "arrow::dataset::DataSourceDiscovery"
 
-    cdef cppclass CFileSystemDataSourceDiscovery "arrow::dataset::FileSystemDataSourceDiscovery"(CDataSourceDiscovery):
+    cdef cppclass CFileSystemDataSourceDiscovery \
+            "arrow::dataset::FileSystemDataSourceDiscovery"(
+                CDataSourceDiscovery):
         @staticmethod
-        CResult[CDataSourceDiscoveryPtr] Make(CFileSystemPtr filesytem,
-                                              CFileStatsVector paths,
-                                              CFileFormatPtr format,
-                                              CFileSystemDiscoveryOptions options)
+        CResult[CDataSourceDiscoveryPtr] Make(
+            CFileSystemPtr filesytem,
+            CFileStatsVector paths,
+            CFileFormatPtr format,
+            CFileSystemDiscoveryOptions options
+        )
         @staticmethod
-        CResult[CDataSourceDiscoveryPtr] Make(CFileSystemPtr filesytem,
-                                              CSelector,
-                                              CFileFormatPtr format,
-                                              CFileSystemDiscoveryOptions options)
+        CResult[CDataSourceDiscoveryPtr] Make(
+            CFileSystemPtr filesytem,
+            CSelector,
+            CFileFormatPtr format,
+            CFileSystemDiscoveryOptions options
+        )

python/pyarrow/tests/test_dataset.py

@@ -16,47 +16,16 @@
 # under the License.
 
 from collections.abc import Generator
-try:
-    import pathlib
-except ImportError:
-    import pathlib2 as pathlib  # py2 compat
-
 import pytest
 
 import pyarrow as pa
-from pyarrow.fs import _MockFileSystem, Selector as FileSelector
+import pyarrow.fs as fs
 
 try:
-    from pyarrow.dataset import (
-        Dataset,
-        DataSource,
-        DataFragment,
-        ParquetFileFormat,
-        SimpleDataFragment,
-        ScanOptions,
-        FileSource,
-        ScanTask,
-        Scanner,
-        ScannerBuilder,
-        ScanOptions,
-        ScanContext,
-        SimpleDataSource,
-        TreeDataSource,
-        FileSystemDataSource,
-        FileSystemDiscoveryOptions,
-        FileSystemDataSourceDiscovery,
-        SchemaPartitionScheme,
-        ScalarExpression,
-        FieldExpression,
-        AndExpression,
-        OrExpression,
-        NotExpression,
-        CompareOperator,
-        ComparisonExpression,
-    )
+    import pyarrow.dataset as ds
 except ImportError as e:
     ds = None
-    raise e
+    raise e  # TODO(kszucs) remove it
 
 # Marks all of the tests in this module
 # Ignore these with pytest ... -m 'not dataset'
@@ -65,10 +34,10 @@ pytestmark = pytest.mark.dataset
 
 @pytest.fixture
 @pytest.mark.parquet
-def fs(table):
+def mockfs(table):
     import pyarrow.parquet as pq
 
-    fs = _MockFileSystem()
+    mockfs = fs._MockFileSystem()
 
     directories = [
         'subdir/1/xxx',
@@ -76,16 +45,16 @@ def fs(table):
     ]
     for i, directory in enumerate(directories):
         path = '{}/file{}.parquet'.format(directory, i)
-        fs.create_dir(directory)
-        with fs.open_output_stream(path) as out:
+        mockfs.create_dir(directory)
+        with mockfs.open_output_stream(path) as out:
             pq.write_table(table, out)
 
-    return fs
+    return mockfs
 
 @pytest.fixture
 def schema():
     return pa.schema([
-        pa.field('i32', pa.int32()),
+        pa.field('i64', pa.int64()),
         pa.field('f64', pa.float64())
     ])
 
@@ -106,22 +75,22 @@ def table(record_batch):
 
 @pytest.fixture
 def simple_data_fragment(record_batch):
-    return SimpleDataFragment([record_batch] * 5)
+    return ds.SimpleDataFragment([record_batch] * 5)
 
 
 @pytest.fixture
 def simple_data_source(simple_data_fragment):
-    return SimpleDataSource([simple_data_fragment] * 4)
+    return ds.SimpleDataSource([simple_data_fragment] * 4)
 
 
 @pytest.fixture
 def tree_data_source(simple_data_source):
-    return TreeDataSource([simple_data_source] * 2)
+    return ds.TreeDataSource([simple_data_source] * 2)
 
 
 @pytest.fixture
-def dataset(simple_data_source, tree_data_source):
-    return Dataset([simple_data_source, tree_data_source])
+def dataset(simple_data_source, tree_data_source, schema):
+    return ds.Dataset([simple_data_source, tree_data_source], schema)
 
 
 def test_scan_options():
@@ -133,22 +102,26 @@ def test_scan_context():
 
 
 def test_simple_data_fragment(record_batch):
-    fragment = SimpleDataFragment([record_batch] * 5)
+    batches = [record_batch] * 5
+    fragment = ds.SimpleDataFragment(batches)
 
-    assert isinstance(fragment, SimpleDataFragment)
+    assert isinstance(fragment, ds.SimpleDataFragment)
     assert fragment.splittable is False
-    assert isinstance(fragment.scan_options, ScanOptions)
+    assert isinstance(fragment.scan_options, ds.ScanOptions)
 
     assert isinstance(fragment.scan(), Generator)
     tasks = list(fragment.scan())
     assert len(tasks) == 5
-    for task in tasks:
-        assert isinstance(task, ScanTask)
+    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 = SimpleDataSource([simple_data_fragment] * 4)
-    assert isinstance(source, SimpleDataSource)
+    source = ds.SimpleDataSource([simple_data_fragment] * 4)
+    assert isinstance(source, ds.SimpleDataSource)
 
     result = source.fragments()
     assert isinstance(result, Generator)
@@ -156,12 +129,12 @@ def test_simple_data_source(simple_data_fragment):
     fragments = list(result)
     assert len(fragments) == 4
     for fragment in fragments:
-        assert isinstance(fragment, DataFragment)
+        assert isinstance(fragment, ds.DataFragment)
 
 
 def test_tree_data_source(simple_data_source):
-    source = TreeDataSource([simple_data_source] * 2)
-    assert isinstance(source, TreeDataSource)
+    source = ds.TreeDataSource([simple_data_source] * 2)
+    assert isinstance(source, ds.TreeDataSource)
 
     result = source.fragments()
     assert isinstance(result, Generator)
@@ -169,25 +142,31 @@ def test_tree_data_source(simple_data_source):
     fragments = list(result)
     assert len(fragments) == 8
     for fragment in fragments:
-        assert isinstance(fragment, DataFragment)
+        assert isinstance(fragment, ds.DataFragment)
 
 
-def test_dataset(simple_data_source, tree_data_source):
-    dataset = Dataset([simple_data_source, tree_data_source])
+def test_dataset(simple_data_source, tree_data_source, schema):
+    dataset = ds.Dataset([simple_data_source, tree_data_source], schema)
 
-    assert isinstance(dataset, Dataset)
+    assert isinstance(dataset, ds.Dataset)
     assert isinstance(dataset.schema, pa.Schema)
     for source in dataset.sources:
-        assert isinstance(source, DataSource)
+        assert isinstance(source, ds.DataSource)
 
-    builder = dataset.new_scan()
-    assert isinstance(builder, ScannerBuilder)
+    condition = ds.ComparisonExpression(
+        ds.CompareOperator.Equal,
+        ds.FieldExpression('i64'),
+        ds.ScalarExpression(1)
+    )
+    builder = dataset.new_scan().use_threads(True).filter(condition)
+    assert isinstance(builder, ds.ScannerBuilder)
+    assert isinstance(builder.schema(), pa.Schema)
 
     scanner = builder.finish()
-    assert isinstance(scanner, Scanner)
+    assert isinstance(scanner, ds.Scanner)
 
     for task in scanner.scan():
-        assert isinstance(task, ScanTask)
+        assert isinstance(task, ds.ScanTask)
         for record_batch in task.scan():
             assert isinstance(record_batch, pa.RecordBatch)
 
@@ -195,12 +174,12 @@ def test_dataset(simple_data_source, tree_data_source):
 def test_scanner(schema, simple_data_source):
     sources = [simple_data_source]
     # FIXME(kszucs): if schema is not set to options it segfaults
-    options = ScanOptions(schema=schema)
-    context = ScanContext()
+    options = ds.ScanOptions(schema=schema)
+    context = ds.ScanContext()
 
-    scanner = Scanner(sources, options, context)
+    scanner = ds.Scanner(sources, options, context)
     for task in scanner.scan():
-        assert isinstance(task, ScanTask)
+        assert isinstance(task, ds.ScanTask)
         for record_batch in task.scan():  # call it execute?
             assert isinstance(record_batch, pa.RecordBatch)
 
@@ -209,79 +188,122 @@ def test_scanner(schema, simple_data_source):
 
 
 def test_scanner_builder(dataset):
-    context = ScanContext()
-    builder = ScannerBuilder(dataset, context)
-    builder.project([])
+    context = ds.ScanContext()
+    builder = ds.ScannerBuilder(dataset, context)
     scanner = builder.finish()
-    assert isinstance(scanner, Scanner)
+    assert isinstance(scanner, ds.Scanner)
+    scanner.scan()
+
+    builder = dataset.new_scan()
+    builder.project(['i64'])
+    scanner = builder.finish()
+    scanner.scan()
+    assert isinstance(scanner, ds.Scanner)
 
 
-def test_projector():
-    pass
-
-
-def test_file_source(fs):
-    source1 = FileSource('/path/to/file.ext', fs, compression=None)
-    source2 = FileSource('/path/to/file.ext.gz', fs, compression='gzip')
+def test_file_source(mockfs):
+    source1 = ds.FileSource('/path/to/file.ext', mockfs, compression=None)
+    source2 = ds.FileSource('/path/to/file.ext.gz', mockfs, compression='gzip')
     assert source1.path == '/path/to/file.ext'
-    assert source1.fs == fs
+    assert source1.fs == mockfs
     assert source1.compression == 0  # None
     assert source2.path == '/path/to/file.ext.gz'
-    assert source2.fs == fs
+    assert source2.fs == mockfs
     assert source2.compression == 2  # 'gzip'
 
 
-def test_file_system_data_source():
-    pass
+@pytest.mark.parametrize('klass', [
+    ds.FileFormat,
+    ds.DataFragment,
+    ds.DataSource,
+    ds.Expression
+])
+def test_abstract_classes(klass):
+    with pytest.raises(TypeError):
+        klass()
 
 
-def test_expression():
-    a = ScalarExpression(1)
-    b = ScalarExpression(1.1)
-    c = ScalarExpression(True)
-    ComparisonExpression(CompareOperator.Equal, a, b)
-    AndExpression(a, b)
-    AndExpression(a, b, c)
-    OrExpression(a, b)
-    OrExpression(a, b, c)
-    NotExpression(OrExpression(a, b, c))
+def test_expression(schema):
+    a = ds.ScalarExpression(1)
+    b = ds.ScalarExpression(1.1)
+    c = ds.ScalarExpression(True)
+
+    equal = ds.ComparisonExpression(ds.CompareOperator.Equal, a, b)
+    assert equal.op() == ds.CompareOperator.Equal
+
+    and_ = ds.AndExpression(a, b)
+    assert isinstance(and_.left_operand, ds.Expression)
+    assert isinstance(and_.right_operand, ds.Expression)
+    assert and_.equals(ds.AndExpression(a, b))
+    assert and_.equals(and_)
+    assert and_ == ds.AndExpression(a, b)
+    assert and_ != 'other object'
+
+    ds.AndExpression(a, b, c)
+    ds.OrExpression(a, b)
+    ds.OrExpression(a, b, c)
+    ds.NotExpression(ds.OrExpression(a, b, c))
+
+    condition = ds.ComparisonExpression(
+        ds.CompareOperator.Greater,
+        ds.FieldExpression('i64'),
+        ds.ScalarExpression(5)
+    )
+    assert condition.validate(schema) == pa.bool_()
+
+    i64_is_5 = ds.ComparisonExpression(
+        ds.CompareOperator.Equal,
+        ds.FieldExpression('i64'),
+        ds.ScalarExpression(5)
+    )
+    i64_is_7 = ds.ComparisonExpression(
+        ds.CompareOperator.Equal,
+        ds.FieldExpression('i64'),
+        ds.ScalarExpression(7)
+    )
+    true_ = ds.ScalarExpression(True)
+    false_ = ds.ScalarExpression(False)
+    assert condition.assume(i64_is_5) == false_
+    assert condition.assume(i64_is_7) == true_
+    assert str(condition) == "(i64 > 5:int64)"
 
 
-def test_file_system_discovery(fs):
-    selector = FileSelector('subdir', recursive=True)
+def test_file_system_discovery(mockfs):
+    selector = fs.Selector('subdir', recursive=True)
     assert selector.base_dir == 'subdir'
     assert selector.recursive is True
 
-    format = ParquetFileFormat()
+    format = ds.ParquetFileFormat()
     assert format.name() == 'parquet'
 
-    options = FileSystemDiscoveryOptions('/')
+    options = ds.FileSystemDiscoveryOptions('/')
     assert options.partition_base_dir == '/'
     assert options.ignore_prefixes == ['.', '_']
     assert options.exclude_invalid_files is True
 
-    discovery = FileSystemDataSourceDiscovery(fs, selector, format, options)
+    discovery = ds.FileSystemDataSourceDiscovery(mockfs, selector, format,
+                                                 options)
     assert isinstance(discovery.inspect(), pa.Schema)
     # assert isinstance(discovery.schema(), pa.Schema)
-    assert isinstance(discovery.finish(), FileSystemDataSource)
+    assert isinstance(discovery.finish(), ds.FileSystemDataSource)
 
-    scheme = SchemaPartitionScheme(
+    scheme = ds.SchemaPartitionScheme(
         pa.schema([
             pa.field('group', pa.int32()),
             pa.field('key', pa.string())
         ])
     )
     discovery.partition_scheme = scheme
-    assert isinstance(discovery.partition_scheme, SchemaPartitionScheme)
+    assert isinstance(discovery.partition_scheme, ds.SchemaPartitionScheme)
 
     data_source = discovery.finish()
-    assert isinstance(data_source, DataSource)
+    assert isinstance(data_source, ds.DataSource)
 
     inspected_schema = discovery.inspect()
-    dataset = Dataset([data_source], inspected_schema)
+    dataset = ds.Dataset([data_source], inspected_schema)
 
     scanner = dataset.new_scan().finish()
     for task in scanner.scan():
-        assert isinstance(task, ScanTask)
+        assert isinstance(task, ds.ScanTask)
         for record_batch in task.scan():
             assert isinstance(record_batch, pa.RecordBatch)