Basic scanning [skip ci]

cpp/src/arrow/dataset/file_base.h

@@ -46,6 +46,7 @@ class Filter;
 /// be read like a file
 class ARROW_DS_EXPORT FileSource {
  public:
+  // NOTE(kszucs): it'd be better to separate the BufferSource from FileSource
   enum SourceType { PATH, BUFFER };
 
   FileSource(std::string path, fs::FileSystem* filesystem,

python/pyarrow/_dataset.pyx

@@ -28,7 +28,7 @@ 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.lib cimport RecordBatch, MemoryPool, Schema, check_status
 from pyarrow.includes.libarrow_dataset cimport *
 
 
@@ -45,77 +45,74 @@ __all__ = [
 ]
 
 
-cdef class ScanOptions:
+cdef class WriteOptions:
+    pass
+
+
+cdef class FileScanOptions(ScanOptions):
+    pass
+
+
+cdef class CsvScanOptions(FileScanOptions):
+    pass
+
+
+cdef class FeatherScanOptions(FileScanOptions):
+    pass
+
+
+cdef class ParquetScanOptions(FileScanOptions):
+    pass
+
+
+cdef class JsonScanOptions(FileScanOptions):
+    pass
+
+
+cdef class FileWriteOptions(WriteOptions):
+    pass
+
+
+cdef class CsvWriteOptions(FileWriteOptions):
+    pass
+
+
+cdef class FeatherWriterOptions(FileWriteOptions):
+    pass
+
+
+cdef class JsonWriterOptions(FileWriteOptions):
+    pass
+
+
+cdef class ParquetWriterOptions(FileWriteOptions):
+    pass
+
+
+cdef class FileFormat:
+    pass
+
+
+cdef class CsvFileFormat(FileFormat):
+    pass
+
+
+cdef class FeatherFileFormat(FileFormat):
+    pass
+
+
+cdef class JsonFileFormat(FileFormat):
+    pass
+
+
+cdef class ParquetFileFormat(FileFormat):
+    pass
+
+
+cdef class FileSource:
 
     cdef:
-        shared_ptr[CScanOptions] wrapped
-        CScanOptions* options
-
-    def __init__(self):
-        cdef shared_ptr[CScanOptions] options
-        self.init(options)
-
-    cdef init(self, const shared_ptr[CScanOptions]& sp):
-        self.wrapped = sp
-        self.options = sp.get()
-
-    @staticmethod
-    cdef ScanOptions wrap(const shared_ptr[CScanOptions]& sp):
-        cdef ScanOptions self = ScanOptions.__new__(ScanOptions)
-        self.init(sp)
-        return self
-
-
-cdef class ScanContext:
-
-    cdef:
-        shared_ptr[CScanContext] wrapped
-        CScanContext* context
-
-    def __init__(self, MemoryPool memory_pool=None):
-        cdef shared_ptr[CScanContext] context
-
-        if memory_pool is not None:
-            context.get().pool = memory_pool.pool
-
-        self.init(context)
-
-    cdef init(self, const shared_ptr[CScanContext]& sp):
-        self.wrapped = sp
-        self.context = sp.get()
-
-    @staticmethod
-    cdef ScanContext wrap(const shared_ptr[CScanContext]& sp):
-        cdef ScanContext self = ScanContext.__new__(ScanContext)
-        self.init(sp)
-        return self
-
-
-cdef class ScanTask:
-
-    cdef:
-        unique_ptr[CScanTask] wrapped
-        CScanTask* task
-
-    cdef init(self, unique_ptr[CScanTask]& sp):
-        self.wrapped = move(sp)
-        self.task = self.wrapped.get()
-
-    @staticmethod
-    cdef ScanTask wrap(unique_ptr[CScanTask]& sp):
-        cdef ScanTask self = ScanTask.__new__(ScanTask)
-        self.init(sp)
-        return self
-
-
-cdef class SimpleScanTask(ScanTask):
-
-    cdef:
-        CSimpleScanTask* simple_task
-
-    cdef init(self, unique_ptr[CScanTask]& sp):
-        ScanTask.init(self, sp)
-        self.simple_task = <CSimpleScanTask*> sp.get()
+        shared_ptr[CFileSource] wrapped
 
 
 cdef class DataFragment:
@@ -133,22 +130,20 @@ cdef class DataFragment:
         self.init(sp)
         return self
 
-    def tasks(self, ScanContext context=None):
+    def scan(self, ScanContext context=None):
         cdef:
-            unique_ptr[CScanTaskIterator] task_iterator
+            CScanTaskIterator task_iterator
             unique_ptr[CScanTask] task
 
         context = context or ScanContext()
-        # check_status(self.fragment.Scan(context.wrapped, &task_iterator))
+        check_status(self.fragment.Scan(context.wrapped, &task_iterator))
 
-        # while True:
-        #     task_iterator.get().Next(&task)
-        #     if task == nullptr:
-        #         raise StopIteration()
-        #     else:
-        #         yield ScanTask.wrap(task)
-
-    scan = tasks
+        while True:
+            task_iterator.Next(&task)
+            if task == nullptr:
+                raise StopIteration()
+            else:
+                yield ScanTask.wrap(task)
 
     @property
     def splittable(self):
@@ -181,6 +176,14 @@ cdef class SimpleDataFragment(DataFragment):
         self.simple_fragment = <CSimpleDataFragment*> sp.get()
 
 
+cdef class FileDataFragment(DataFragment):
+    pass
+
+
+cdef class ParquetDataFragment(FileDataFragment):
+    pass
+
+
 cdef class DataSource:
 
     cdef:
@@ -191,26 +194,50 @@ cdef class DataSource:
         self.wrapped = sp
         self.source = sp.get()
 
+    @staticmethod
+    cdef wrap(shared_ptr[CDataSource]& sp):
+        cdef DataSource self
+
+        typ = frombytes(sp.get().type())
+        if typ == 'tree_data_source':
+            self = TreeDataSource.__new__(TreeDataSource)
+        elif typ == 'simple_data_source':
+            self = SimpleDataSource.__new__(SimpleDataSource)
+        elif typ == 'filesystem_data_source':
+            self = FileSystemDataSource.__new__(FileSystemDataSource)
+        else:
+            self = DataSource.__new__(DataSource)
+
+        self.init(sp)
+        return self
+
+    cdef shared_ptr[CDataSource] unwrap(self):
+        return self.wrapped
+
     @property
     def type(self):
         return frombytes(self.source.type())
 
+    # def __iter__(self):
+    #     for fragment in self.fragments():
+    #         yield fragment
+
     def fragments(self, ScanOptions options=None):
         cdef:
-            unique_ptr[CDataFragmentIterator] fragment_iterator
+            CDataFragmentIterator fragment_iterator
             shared_ptr[CDataFragment] fragment
 
         options = options or ScanOptions()
 
         # NOTE(kszucs): shouldn't we return Status here?
-        # fragment_iterator = self.source.GetFragments(options.wrapped)
+        fragment_iterator = self.source.GetFragments(options.wrapped)
 
-        # while True:
-        #     fragment_iterator.get().Next(&fragment)
-        #     if fragment == nullptr:
-        #         raise StopIteration()
-        #     else:
-        #         yield DataFragment.wrap(fragment)
+        while True:
+            fragment_iterator.Next(&fragment)
+            if fragment == nullptr:
+                raise StopIteration()
+            else:
+                yield DataFragment.wrap(fragment)
 
     get_fragments = fragments
 
@@ -237,47 +264,250 @@ cdef class SimpleDataSource(DataSource):
         self.simple_source = <CSimpleDataSource*> sp.get()
 
 
+cdef class TreeDataSource(DataSource):
+
+    cdef:
+        CTreeDataSource* tree_source
+
+    def __init__(self, data_sources):
+        cdef:
+            DataSource child
+            vector[shared_ptr[CDataSource]] children
+            shared_ptr[CTreeDataSource] tree_source
+
+        for child in data_sources:
+            children.push_back(child.wrapped)
+
+        tree_source = make_shared[CTreeDataSource](children)
+        self.init(<shared_ptr[CDataSource]> tree_source)
+
+    cdef void init(self, const shared_ptr[CDataSource]& sp):
+        DataSource.init(self, sp)
+        self.tree_source = <CTreeDataSource*> sp.get()
+
+
+cdef class FileSystemDataSource(DataSource):
+
+    cdef:
+        CFileSystemDataSource* filesystem_source
+
+    def __init__(self, filesystem, file_stats, source_partition,
+                 path_partitions, file_format):
+        cdef:
+            DataSource child
+            vector[shared_ptr[CDataSource]] children
+            shared_ptr[CFileSystemDataSource] filesystem_source
+
+        # TODO(kszucs)
+        # filesystem_source = CFileSystemDataSource.Make(
+        #     fs::FileSystem* filesystem, std::vector<fs::FileStats> stats,
+        #     std::shared_ptr<Expression> source_partition,
+        #     PathPartitions partitions, std::shared_ptr<FileFormat> format,
+        #     std::shared_ptr<DataSource>* out
+        # )
+        self.init(<shared_ptr[CDataSource]> filesystem_source)
+
+    cdef void init(self, const shared_ptr[CDataSource]& sp):
+        DataSource.init(self, sp)
+        self.filesystem_source = <CFileSystemDataSource*> sp.get()
+
+
 cdef class Dataset:
 
     cdef:
         shared_ptr[CDataset] wrapped
         CDataset* dataset
 
-    # def __init__(self, data_sources, Schema schema=None):
-    #     cdef:
-    #         DataSource source
-    #         vector[shared_ptr[CDataSource]] sources
-    #         shared_ptr[CDataset] dataset
+    def __init__(self, data_sources, Schema schema=None):
+        cdef:
+            DataSource source
+            vector[shared_ptr[CDataSource]] sources
+            shared_ptr[CDataset] dataset
 
-    #     for source in data_sources:
-    #         sources.push_back(source.wrapped)
+        for source in data_sources:
+            sources.push_back(source.wrapped)
 
-    #     if schema is None:
-    #         dataset = make_shared[CDataset](sources)
-    #     else:
-    #         dataset = make_shared[CDataset](sources, schema.sp_schema)
+        check_status(CDataset.Make(sources, schema.sp_schema, &dataset))
+        # if schema is None:
+        #     dataset = make_shared[CDataset](sources)
+        # else:
+        #     dataset = make_shared[CDataset](sources, schema.sp_schema)
+        self.init(dataset)
 
-    #     self.init(dataset)
-
-    # cdef void init(self, const shared_ptr[CDataset]& sp):
-    #     self.wrapped = sp
-    #     self.dataset = sp.get()
+    cdef void init(self, const shared_ptr[CDataset]& sp):
+        self.wrapped = sp
+        self.dataset = sp.get()
 
     def new_scan(self):
-        pass
+        cdef:
+            unique_ptr[CScannerBuilder] builder
+            unique_ptr[CScanner] scanner
+
+        # TODO(kszucs): add other options like filter, schema etc.
+        check_status(self.dataset.NewScan(&builder))
+        check_status(builder.get().Finish(&scanner))
+
+        return Scanner.wrap(scanner)
 
     @property
     def sources(self):
-        pass
+        cdef vector[shared_ptr[CDataSource]] sources = self.dataset.sources()
+        return [DataSource.wrap(source) for source in sources]
 
     @property
     def schema(self):
-        pass
+        return pyarrow_wrap_schema(self.dataset.schema())
 
-    # def infer_schema(self):
-    #     pass
 
-    # def replace_schema(self, Schema new_schema):
-    #     pass
+cdef class ScanOptions:
 
-    # with_schema = replace_schema
+    cdef:
+        shared_ptr[CScanOptions] wrapped
+        CScanOptions* options
+
+    def __init__(self):
+        cdef shared_ptr[CScanOptions] options
+        self.init(options)
+
+    cdef init(self, const shared_ptr[CScanOptions]& sp):
+        self.wrapped = sp
+        self.options = sp.get()
+
+    @staticmethod
+    cdef ScanOptions 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
+
+
+cdef class ScanContext:
+
+    cdef:
+        shared_ptr[CScanContext] wrapped
+        CScanContext* context
+
+    def __init__(self, MemoryPool memory_pool=None):
+        cdef shared_ptr[CScanContext] context
+
+        if memory_pool is not None:
+            context.get().pool = memory_pool.pool
+
+        self.init(context)
+
+    cdef init(self, const shared_ptr[CScanContext]& sp):
+        self.wrapped = sp
+        self.context = sp.get()
+
+    @staticmethod
+    cdef ScanContext wrap(const shared_ptr[CScanContext]& sp):
+        cdef ScanContext self = ScanContext.__new__(ScanContext)
+        self.init(sp)
+        return self
+
+    cdef inline shared_ptr[CScanContext] unwrap(self):
+        return self.wrapped
+
+
+cdef class ScanTask:
+
+    cdef:
+        unique_ptr[CScanTask] wrapped
+        CScanTask* task
+
+    cdef init(self, unique_ptr[CScanTask]& up):
+        self.wrapped = move(up)
+        self.task = self.wrapped.get()
+
+    @staticmethod
+    cdef wrap(unique_ptr[CScanTask]& up):
+        cdef SimpleScanTask self = SimpleScanTask.__new__(SimpleScanTask)
+        self.init(up)
+        return self
+
+    def scan(self):
+        cdef:
+            CRecordBatchIterator it = self.task.Scan()
+            shared_ptr[CRecordBatch] record_batch
+
+        while True:
+            it.Next(&record_batch)
+            if record_batch == nullptr:
+                raise StopIteration()
+            else:
+                yield pyarrow_wrap_batch(record_batch)
+
+
+cdef class SimpleScanTask(ScanTask):
+
+    cdef:
+        CSimpleScanTask* simple_task
+
+    cdef init(self, unique_ptr[CScanTask]& up):
+        ScanTask.init(self, up)
+        self.simple_task = <CSimpleScanTask*> up.get()
+
+
+cdef class Scanner:
+
+    cdef:
+        unique_ptr[CScanner] wrapped
+        CScanner* scanner
+
+    cdef void init(self, unique_ptr[CScanner]& up):
+        self.wrapped = move(up)
+        self.scanner = self.wrapped.get()
+
+    @staticmethod
+    cdef wrap(unique_ptr[CScanner]& up):
+        cdef Scanner self = SimpleScanner.__new__(SimpleScanner)
+        self.init(up)
+        return self
+
+    def scan(self):
+        cdef:
+            CScanTaskIterator task_iterator = self.scanner.Scan()
+            unique_ptr[CScanTask] task
+
+        while True:
+            task_iterator.Next(&task)
+            if task == nullptr:
+                raise StopIteration()
+            else:
+                yield ScanTask.wrap(task)
+
+    def to_table(self):
+        cdef shared_ptr[CTable] table
+
+        with nogil:
+            check_status(self.scanner.ToTable(&table))
+
+        return pyarrow_wrap_table(table)
+
+
+cdef class SimpleScanner(Scanner):
+
+    cdef:
+        CSimpleScanner* simple_scanner
+
+    def __init__(self, data_sources, ScanOptions options, ScanContext context):
+        cdef:
+            DataSource source
+            CDataSourceVector sources
+            CSimpleScanner* simple_scanner
+            unique_ptr[CScanner] scanner
+
+        for source in data_sources:
+            sources.push_back(source.unwrap())
+
+        simple_scanner = new CSimpleScanner(
+            sources, options.unwrap(), context.unwrap()
+        )
+        scanner.reset(simple_scanner)
+        self.init(scanner)
+
+
+# cdef class ScannerBuilder:
+#     pass

python/pyarrow/dataset.py

@@ -22,4 +22,17 @@ import sys
 if sys.version_info < (3,):
     raise ImportError("Python Dataset bindings require Python 3")
 
-from pyarrow._dataset import *  # noqa
+from pyarrow._dataset import (
+    Dataset,
+    DataSource,
+    DataFragment,
+    SimpleDataFragment,
+    ScanOptions,
+    ScanTask,
+    Scanner,
+    SimpleScanner,
+    ScanOptions,
+    ScanContext,
+    SimpleDataSource,
+    TreeDataSource
+)

python/pyarrow/includes/common.pxd

@@ -19,7 +19,7 @@
 
 from libc.stdint cimport *
 from libcpp cimport bool as c_bool, nullptr
-from libcpp.memory cimport shared_ptr, unique_ptr, make_shared
+from libcpp.memory cimport shared_ptr, unique_ptr, make_shared, make_unique
 from libcpp.string cimport string as c_string
 from libcpp.utility cimport pair
 from libcpp.vector cimport vector

python/pyarrow/includes/libarrow_dataset.pxd

@@ -23,6 +23,13 @@ from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport *
 
 
+cdef extern from "arrow/api.h" namespace "arrow" nogil:
+
+    cdef cppclass CRecordBatchIterator "arrow::RecordBatchIterator"(
+            CIterator[shared_ptr[CRecordBatch]]):
+        pass
+
+
 cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
 
     # cdef enum CFilter" arrow::dataset::Filter::type":
@@ -36,12 +43,20 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     cdef cppclass CDataFragmentVector "arrow::dataset::DataFragmentVector":
         pass
 
-    cdef cppclass CDataSourceVector "arrow::dataset::DataSourceVector":
+    cdef cppclass CDataSourceVector "arrow::dataset::DataSourceVector"(
+            vector[shared_ptr[CDataSource]]):
         pass
 
     cdef cppclass CWriteOptions "arrow::dataset::WriteOptions":
         pass
 
+    cdef cppclass CExpression "arrow::dataset::Expression":
+        # CExpression(ExpressionType::type type) : type_(type) {}
+        c_bool Equals(const CExpression& other) const
+        c_bool Equals(const shared_ptr[CExpression]& other) const
+        c_bool IsNull() const
+        CStatus Validate(const CSchema& schema, shared_ptr[CDataType]* out) const
+
     cdef cppclass CFilter "arrow::dataset::Filter":
         pass
 
@@ -53,32 +68,43 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CMemoryPool* pool
 
     cdef cppclass CScanTask" arrow::dataset::ScanTask":
-        pass
+        CRecordBatchIterator Scan()
 
-    cdef cppclass CSimpleScanTask" arrow::dataset::SimpleScanTask":
+    cdef cppclass CSimpleScanTask "arrow::dataset::SimpleScanTask":
         pass
 
     cdef cppclass CScanTaskIterator "arrow::dataset::ScanTaskIterator"(
             CIterator[unique_ptr[CScanTask]]):
         pass
 
-    cdef cppclass CScannerBuilder" arrow::dataset::ScannerBuilder":
+    cdef cppclass CScannerBuilder "arrow::dataset::ScannerBuilder":
         CScannerBuilder(shared_ptr[CDataset] dataset,
                         shared_ptr[CScanContext] scan_context)
-        CScannerBuilder* Project(const vector[c_string]& columns)
-        CScannerBuilder* AddFilter(const shared_ptr[CFilter]& filter)
-        CScannerBuilder* SetGlobalFileOptions(
-            shared_ptr[CFileScanOptions] options)
-        CScannerBuilder* IncludePartitionKeys(c_bool)
-        unique_ptr[CScanner] Finish()
+        CStatus Project(const vector[c_string]& columns)
+        CStatus Filter(const CExpression& filter)
+        CStatus Filter(shared_ptr[CExpression] filter)
+        CStatus UseThreads(c_bool use_threads)
+        CStatus Finish(unique_ptr[CScanner]* out) const
+        shared_ptr[CSchema] schema() const
 
-    cdef cppclass CScanTaskIterator" arrow::dataset::ScanTaskIterator":
+    cdef cppclass CScanTaskIterator "arrow::dataset::ScanTaskIterator":
         pass
 
-    cdef cppclass CScanner" arrow::dataset::Scanner":
-        unique_ptr[CScanTaskIterator] Scan()
+    cdef cppclass CScanner "arrow::dataset::Scanner":
+        c_bool use_threads
+        shared_ptr[CExpression] filter
+        # shared_ptr[CExpressionEvaluator] evaluator
+        shared_ptr[CSchema] schema
+        # shared_ptr[CRecordBatchProjector] projector
+        CScanTaskIterator Scan()
+        CStatus ToTable(shared_ptr[CTable]* out)
 
-    cdef cppclass CDataFragment" arrow::dataset::DataFragment":
+    cdef cppclass CSimpleScanner "arrow::dataset::SimpleScanner"(CScanner):
+        CSimpleScanner(CDataSourceVector sources,
+                       shared_ptr[CScanOptions] options,
+                       shared_ptr[CScanContext] context)
+
+    cdef cppclass CDataFragment "arrow::dataset::DataFragment":
         CStatus Scan(shared_ptr[CScanContext] scan_context,
                      CScanTaskIterator* out)
         c_bool splittable()
@@ -103,6 +129,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         pass
 
     cdef cppclass CDataset" arrow::dataset::Dataset":
+        @staticmethod
         CStatus Make(vector[shared_ptr[CDataSource]] sources,
                      shared_ptr[CSchema] schema, shared_ptr[CDataset] *out)
         CStatus NewScan(unique_ptr[CScannerBuilder]* out)
@@ -115,7 +142,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
             CScanOptions):
         c_string file_type()
 
-    cdef cppclass CFileSource" arrow::dataset::FileSource":
+    cdef cppclass CFileSource "arrow::dataset::FileSource":
         # enum SourceType { PATH, BUFFER };
         CFileSource(c_string path, CFileSystem filesystem,
                     CompressionType compression)
@@ -132,7 +159,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
             CWriteOptions):
         c_string file_type()
 
-    cdef cppclass CFileFormat" arrow::dataset::FileFormat":
+    cdef cppclass CFileFormat "arrow::dataset::FileFormat":
         c_string name()
         c_bool IsKnownExtension(const c_string& ext)
         CStatus ScanFile(const CFileSource& source,
@@ -154,14 +181,15 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         shared_ptr[CFileFormat] format()
         shared_ptr[CScanOptions] scan_options()
 
-    cdef cppclass CFileSystemBasedDataSource \
+    # NOTE(kszucs): removed Based from the class name
+    cdef cppclass CFileSystemDataSource \
             "arrow::dataset::FileSystemBasedDataSource"(CDataSource):
         @staticmethod
         CStatus Make(CFileSystem* filesystem,
                      const CSelector& selector,
                      shared_ptr[CFileFormat] format,
                      shared_ptr[CScanOptions] scan_options,
-                     unique_ptr[CFileSystemBasedDataSource]* out)
+                     unique_ptr[CFileSystemDataSource]* out)
         c_string type()
         unique_ptr[CDataFragmentIterator] GetFragments(
             shared_ptr[CScanOptions] options)

python/pyarrow/tests/test_dataset.py

@@ -23,11 +23,18 @@ import pyarrow as pa
 
 try:
     from pyarrow.dataset import (
+        Dataset,
+        DataSource,
         DataFragment,
         SimpleDataFragment,
         ScanOptions,
         ScanTask,
-        SimpleDataSource
+        Scanner,
+        ScanOptions,
+        ScanContext,
+        SimpleScanner,
+        SimpleDataSource,
+        TreeDataSource
     )
 except ImportError as e:
     ds = None
@@ -39,11 +46,15 @@ pytestmark = pytest.mark.dataset
 
 
 @pytest.fixture
-def record_batch():
-    schema = pa.schema([
+def schema():
+    return pa.schema([
         pa.field('i32', pa.int32()),
         pa.field('f64', pa.float64())
     ])
+
+
+@pytest.fixture
+def record_batch(schema):
     data = [
         list(range(5)),
         list(map(float, range(5)))
@@ -51,6 +62,21 @@ def record_batch():
     return pa.record_batch(data, schema=schema)
 
 
+@pytest.fixture
+def simple_data_fragment(record_batch):
+    return SimpleDataFragment([record_batch] * 5)
+
+
+@pytest.fixture
+def simple_data_source(simple_data_fragment):
+    return SimpleDataSource([simple_data_fragment] * 4)
+
+
+@pytest.fixture
+def tree_data_source(simple_data_source):
+    return TreeDataSource([simple_data_source] * 2)
+
+
 def test_scan_options():
     pass
 
@@ -66,20 +92,15 @@ def test_simple_data_fragment(record_batch):
     assert fragment.splittable is False
     assert isinstance(fragment.scan_options, ScanOptions)
 
-    # scan is an alias for tasks
-    for result in [fragment.scan(), fragment.tasks()]:
-        result = fragment.tasks()
-        assert isinstance(result, Generator)
-
-        tasks = list(result)
-        assert len(tasks) == 5
-        for task in tasks:
-            assert isinstance(task, ScanTask)
+    assert isinstance(fragment.scan(), Generator)
+    tasks = list(fragment.scan())
+    assert len(tasks) == 5
+    for task in tasks:
+        assert isinstance(task, ScanTask)
 
 
-def test_simple_data_source(record_batch):
-    fragment = SimpleDataFragment([record_batch] * 5)
-    source = SimpleDataSource([fragment] * 4)
+def test_simple_data_source(simple_data_fragment):
+    source = SimpleDataSource([simple_data_fragment] * 4)
 
     assert isinstance(source, SimpleDataSource)
     assert source.type == 'simple_data_source'
@@ -91,3 +112,50 @@ def test_simple_data_source(record_batch):
     assert len(fragments) == 4
     for fragment in fragments:
         assert isinstance(fragment, DataFragment)
+
+
+def test_tree_data_source(simple_data_source):
+    source = TreeDataSource([simple_data_source] * 2)
+
+    assert isinstance(source, TreeDataSource)
+    assert source.type == 'tree_data_source'
+
+    result = source.fragments()
+    assert isinstance(result, Generator)
+
+    fragments = list(result)
+    assert len(fragments) == 8
+    for fragment in fragments:
+        assert isinstance(fragment, DataFragment)
+
+
+def test_dataset(simple_data_source, tree_data_source):
+    dataset = Dataset([simple_data_source, tree_data_source])
+
+    assert isinstance(dataset, Dataset)
+    assert isinstance(dataset.schema, pa.Schema)
+    for source in dataset.sources:
+        assert isinstance(source, DataSource)
+
+    scanner = dataset.new_scan()
+    assert isinstance(scanner, Scanner)
+    assert isinstance(scanner, SimpleScanner)
+
+    for task in scanner.scan():
+        assert isinstance(task, ScanTask)
+        for record_batch in task.scan():
+            assert isinstance(record_batch, pa.RecordBatch)
+
+
+# def test_scanner(simple_data_source):
+#     sources = [simple_data_source]
+#     options = ScanOptions()
+#     context = ScanContext()
+#     scanner = SimpleScanner(sources, options, context)
+
+#     for task in scanner.scan():
+#         assert isinstance(task, ScanTask)
+
+
+def test_projector():
+    pass