data fragments and data sources

cpp/src/arrow/dataset/dataset.h

@@ -123,6 +123,8 @@ class ARROW_DS_EXPORT SimpleDataSource : public DataSource {
 
   DataFragmentIterator GetFragmentsImpl(ScanOptionsPtr options) override;
 
+  // NOTE(kszucs): shouldn't we use an enum here and use name() instead?
+  // similarly like the DataType API works
   std::string type() const override { return "simple_data_source"; }
 
  private:

cpp/src/arrow/dataset/scanner.h

@@ -85,6 +85,7 @@ class ARROW_DS_EXPORT ScanTask {
   /// \brief Iterate through sequence of materialized record batches
   /// resulting from the Scan. Execution semantics encapsulated in the
   /// particular ScanTask implementation
+  // TODO(kszucs): call it execute?
   virtual Result<RecordBatchIterator> Scan() = 0;
 
   virtual ~ScanTask() = default;

python/pyarrow/_dataset.pyx

@@ -0,0 +1,283 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# cython: language_level = 3
+
+from __future__ import absolute_import
+
+import collections
+import six
+
+from cython.operator cimport dereference as deref
+
+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.includes.libarrow_dataset cimport *
+
+
+__all__ = [
+    'ScanOptions',
+    'ScanContext',
+    'ScanTask',
+    'SimpleScanTask',
+    'DataFragment',
+    'SimpleDataFragment',
+    'DataSource',
+    'SimpleDataSource',
+    'Dataset'
+]
+
+
+cdef class ScanOptions:
+
+    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()
+
+
+cdef class DataFragment:
+    cdef:
+        shared_ptr[CDataFragment] wrapped
+        CDataFragment* fragment
+
+    cdef void init(self, const shared_ptr[CDataFragment]& sp):
+        self.wrapped = sp
+        self.fragment = sp.get()
+
+    @staticmethod
+    cdef DataFragment wrap(shared_ptr[CDataFragment]& sp):
+        cdef DataFragment self = DataFragment.__new__(DataFragment)
+        self.init(sp)
+        return self
+
+    def tasks(self, ScanContext context=None):
+        cdef:
+            unique_ptr[CScanTaskIterator] task_iterator
+            unique_ptr[CScanTask] task
+
+        context = context or ScanContext()
+        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
+
+    @property
+    def splittable(self):
+        return self.fragment.splittable()
+
+    @property
+    def scan_options(self):
+        return ScanOptions.wrap(self.fragment.scan_options())
+
+
+cdef class SimpleDataFragment(DataFragment):
+
+    cdef:
+        CSimpleDataFragment* simple_fragment
+
+    def __init__(self, record_batches):
+        cdef:
+            RecordBatch batch
+            vector[shared_ptr[CRecordBatch]] batches
+            shared_ptr[CSimpleDataFragment] simple_fragment
+
+        for batch in record_batches:
+            batches.push_back(batch.sp_batch)
+
+        simple_fragment = make_shared[CSimpleDataFragment](batches)
+        self.init(<shared_ptr[CDataFragment]> simple_fragment)
+
+    cdef void init(self, const shared_ptr[CDataFragment]& sp):
+        DataFragment.init(self, sp)
+        self.simple_fragment = <CSimpleDataFragment*> sp.get()
+
+
+cdef class DataSource:
+
+    cdef:
+        shared_ptr[CDataSource] wrapped
+        CDataSource* source
+
+    cdef void init(self, const shared_ptr[CDataSource]& sp):
+        self.wrapped = sp
+        self.source = sp.get()
+
+    @property
+    def type(self):
+        return frombytes(self.source.type())
+
+    def fragments(self, ScanOptions options=None):
+        cdef:
+            unique_ptr[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)
+
+        while True:
+            fragment_iterator.get().Next(&fragment)
+            if fragment == nullptr:
+                raise StopIteration()
+            else:
+                yield DataFragment.wrap(fragment)
+
+    get_fragments = fragments
+
+
+cdef class SimpleDataSource(DataSource):
+
+    cdef:
+        CSimpleDataSource* simple_source
+
+    def __init__(self, data_fragments):
+        cdef:
+            DataFragment fragment
+            vector[shared_ptr[CDataFragment]] fragments
+            shared_ptr[CSimpleDataSource] simple_source
+
+        for fragment in data_fragments:
+            fragments.push_back(fragment.wrapped)
+
+        simple_source = make_shared[CSimpleDataSource](fragments)
+        self.init(<shared_ptr[CDataSource]> simple_source)
+
+    cdef void init(self, const shared_ptr[CDataSource]& sp):
+        DataSource.init(self, sp)
+        self.simple_source = <CSimpleDataSource*> 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
+#
+#         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)
+#
+#         self.init(dataset)
+#
+#     cdef void init(self, const shared_ptr[CDataset]& sp):
+#         self.wrapped = sp
+#         self.dataset = sp.get()
+#
+#     def new_scan(self):
+#         pass
+#
+#     @property
+#     def sources(self):
+#         pass
+#
+#     @property
+#     def schema(self):
+#         pass
+#
+#     def infer_schema(self):
+#         pass
+#
+#     def replace_schema(self, Schema new_schema):
+#         pass
+#
+#     with_schema = replace_schema

python/pyarrow/dataset.py

@@ -22,4 +22,4 @@ import sys
 if sys.version_info < (3,):
     raise ImportError("Python Dataset bindings require Python 3")
 
-from pyarrow._flight import *  # noqa
+from pyarrow._dataset import *  # noqa

python/pyarrow/includes/common.pxd

@@ -33,6 +33,24 @@ cimport cpython
 cdef extern from * namespace "std" nogil:
     cdef shared_ptr[T] static_pointer_cast[T, U](shared_ptr[U])
 
+# vendored from the cymove project https://github.com/ozars/cymove
+cdef extern from * namespace "cymove":
+    """
+    #include <type_traits>
+    #include <utility>
+    namespace cymove {
+    template <typename T>
+    inline typename std::remove_reference<T>::type&& cymove(T& t) {
+        return std::move(t);
+    }
+    template <typename T>
+    inline typename std::remove_reference<T>::type&& cymove(T&& t) {
+        return std::move(t);
+    }
+    }  // namespace cymove
+    """
+    cdef T move" cymove::cymove"[T](T)
+
 cdef extern from "arrow/python/platform.h":
     pass
 

python/pyarrow/includes/libarrow.pxd

@@ -1668,6 +1668,12 @@ cdef extern from 'arrow/util/compression.h' namespace 'arrow' nogil:
         int64_t MaxCompressedLen(int64_t input_len, const uint8_t* input)
 
 
+cdef extern from 'arrow/util/iterator.h' namespace 'arrow' nogil:
+    cdef cppclass CIterator" arrow::Iterator"[T]:
+        CStatus Next(T* out)
+        CStatus Visit[Visitor](Visitor&& visitor)
+
+
 cdef extern from 'arrow/util/thread_pool.h' namespace 'arrow' nogil:
     int GetCpuThreadPoolCapacity()
     CStatus SetCpuThreadPoolCapacity(int threads)

python/pyarrow/includes/libarrow_dataset.pxd

@@ -36,12 +36,19 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         pass
 
     cdef cppclass CScanOptions" arrow::dataset::ScanOptions":
+        # const shared_ptr[DataSelector]& selector()
         pass
 
     cdef cppclass CScanContext" arrow::dataset::ScanContext":
+        CMemoryPool* pool
+
+    cdef cppclass CScanTask" arrow::dataset::ScanTask":
         pass
 
-    cdef cppclass CScanTaskIterator" arrow::dataset::ScanTaskIterator":
+    cdef cppclass CSimpleScanTask" arrow::dataset::SimpleScanTask":
+        pass
+
+    cdef cppclass CScanTaskIterator" arrow::dataset::ScanTaskIterator"(CIterator[unique_ptr[CScanTask]]):
         pass
 
     cdef cppclass CScannerBuilder" arrow::dataset::ScannerBuilder":
@@ -60,6 +67,9 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     cdef cppclass CScanner" arrow::dataset::Scanner":
         unique_ptr[CScanTaskIterator] Scan()
 
+    cdef cppclass CDataFragmentIterator" arrow::dataset::DataFragmentIterator"(CIterator[shared_ptr[CDataFragment]]):
+        pass
+
     cdef cppclass CDataFragment" arrow::dataset::DataFragment":
         CStatus Scan(shared_ptr[CScanContext] scan_context,
                      unique_ptr[CScanTaskIterator]* out)
@@ -68,10 +78,6 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
 
     cdef cppclass CSimpleDataFragment" arrow::dataset::SimpleDataFragment"(CDataFragment):
         CSimpleDataFragment(vector[shared_ptr[CRecordBatch]] record_batches)
-        CStatus Scan(shared_ptr[CScanContext] scan_context,
-                     unique_ptr[CScanTaskIterator]* out)
-        c_bool splittable()
-        shared_ptr[CScanOptions] scan_options()
 
     cdef cppclass CDataFragmentIterator" arrow::dataset::DataFragmentIterator":
         pass
@@ -91,10 +97,10 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         c_string type()
 
     cdef cppclass CDataset" arrow::dataset::Dataset":
-        CDataset(shared_ptr[CDataSource] source, shared_ptr[CSchema] schema)
+        CDataset(shared_ptr[CDataSource] source,
+                 shared_ptr[CSchema] schema = NULLPTR)
         CDataset(const vector[shared_ptr[CDataSource]]& sources,
-                 shared_ptr[CSchema] schema)
-
+                 shared_ptr[CSchema] schema = NULLPTR)
         CScannerBuilder NewScan()
         vector[shared_ptr[CDataSource]]& sources()
         shared_ptr[CSchema] schema()

python/pyarrow/tests/conftest.py

@@ -41,7 +41,6 @@ h.settings.register_profile('debug', max_examples=10,
 # examples try: pytest pyarrow -sv --only-hypothesis --hypothesis-profile=debug
 h.settings.load_profile(os.environ.get('HYPOTHESIS_PROFILE', 'dev'))
 
-
 groups = [
     'cython',
     'dataset',
@@ -62,7 +61,6 @@ groups = [
     'requires_testing_data',
 ]
 
-
 defaults = {
     'cython': False,
     'dataset': False,

python/pyarrow/tests/test_dataset.py

@@ -0,0 +1,93 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from collections.abc import Generator
+
+import pytest
+
+import pyarrow as pa
+
+try:
+    from pyarrow.dataset import (
+        DataFragment,
+        SimpleDataFragment,
+        ScanOptions,
+        ScanTask,
+        SimpleDataSource
+    )
+except ImportError as e:
+    ds = None
+    raise e
+
+# Marks all of the tests in this module
+# Ignore these with pytest ... -m 'not dataset'
+pytestmark = pytest.mark.dataset
+
+
+@pytest.fixture
+def record_batch():
+    schema = pa.schema([
+        pa.field('i32', pa.int32()),
+        pa.field('f64', pa.float64())
+    ])
+    data = [
+        list(range(5)),
+        list(map(float, range(5)))
+    ]
+    return pa.record_batch(data, schema=schema)
+
+
+def test_scan_options():
+    pass
+
+
+def test_scan_context():
+    pass
+
+
+def test_simple_data_fragment(record_batch):
+    fragment = SimpleDataFragment([record_batch] * 5)
+
+    assert isinstance(fragment, SimpleDataFragment)
+    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)
+
+
+def test_simple_data_source(record_batch):
+    fragment = SimpleDataFragment([record_batch] * 5)
+    source = SimpleDataSource([fragment] * 4)
+
+    assert isinstance(source, SimpleDataSource)
+    assert source.type == 'simple_data_source'
+
+    result = source.fragments()
+    assert isinstance(result, Generator)
+
+    fragments = list(result)
+    assert len(fragments) == 4
+    for fragment in fragments:
+        assert isinstance(fragment, DataFragment)