Downcast data fragments

cpp/src/arrow/dataset/dataset.h

@@ -44,6 +44,8 @@ class ARROW_DS_EXPORT DataFragment {
   /// scanning
   virtual bool splittable() const = 0;
 
+  virtual std::string type() const = 0;
+
   /// \brief Filtering, schema reconciliation, and partition options to use when
   /// scanning this fragment. May be nullptr, which indicates that no filtering
   /// or schema reconciliation will be performed and all partitions will be
@@ -76,6 +78,8 @@ class ARROW_DS_EXPORT SimpleDataFragment : public DataFragment {
   SimpleDataFragment(std::vector<std::shared_ptr<RecordBatch>> record_batches,
                      ScanOptionsPtr scan_options);
 
+  std::string type() const override { return "simple_data_fragment"; }
+
   Result<ScanTaskIterator> Scan(ScanContextPtr context) override;
 
   bool splittable() const override { return false; }
@@ -123,8 +127,6 @@ 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/file_base.h

@@ -158,6 +158,8 @@ class ARROW_DS_EXPORT FileDataFragment : public DataFragment {
         source_(source),
         format_(std::move(format)) {}
 
+  std::string type() const override { return "file_data_fragment"; }
+
   Result<ScanTaskIterator> Scan(ScanContextPtr context) override;
 
   const FileSource& source() const { return source_; }

cpp/src/arrow/dataset/file_parquet.h

@@ -70,6 +70,8 @@ class ARROW_DS_EXPORT ParquetFragment : public FileDataFragment {
   ParquetFragment(const FileSource& source, ScanOptionsPtr options)
       : FileDataFragment(source, std::make_shared<ParquetFileFormat>(), options) {}
 
+  std::string type() const override { return "parquet_data_fragment"; }
+
   bool splittable() const override { return true; }
 };
 

python/pyarrow/_dataset.pyx

@@ -38,26 +38,6 @@ def _forbid_instantiation(klass):
     )
 
 
-cdef class WriteOptions:
-    pass
-
-
-cdef class FileScanOptions(ScanOptions):
-    pass
-
-
-cdef class ParquetScanOptions(FileScanOptions):
-    pass
-
-
-cdef class FileWriteOptions(WriteOptions):
-    pass
-
-
-cdef class ParquetWriterOptions(FileWriteOptions):
-    pass
-
-
 cdef class FileFormat:
 
     cdef:
@@ -368,7 +348,18 @@ cdef class DataFragment:
 
     @staticmethod
     cdef wrap(shared_ptr[CDataFragment]& sp):
-        cdef DataFragment self = DataFragment.__new__(DataFragment)
+        cdef DataFragment self
+
+        typ = frombytes(sp.get().type())
+        if typ == 'simple_data_fragment':
+            self = SimpleDataFragment.__new__(SimpleDataFragment)
+        elif typ == 'file_data_fragment':
+            self = FileDataFragment.__new__(FileDataFragment)
+        elif typ == 'parquet_data_fragment':
+            self = ParquetDataFragment.__new__(ParquetDataFragment)
+        else:
+            raise TypeError(typ)
+
         self.init(sp)
         return self
 
@@ -622,10 +613,9 @@ cdef class ScanOptions:
         CScanOptions* options
 
     def __init__(self, Schema schema=None):
-        cdef shared_ptr[CScanOptions] options = CScanOptions.Defaults()
+        self.init(CScanOptions.Defaults())
         if schema is not None:
-            options.get().schema = pyarrow_unwrap_schema(schema)
-        self.init(options)
+            self.schema = schema
 
     cdef init(self, const shared_ptr[CScanOptions]& sp):
         self.wrapped = sp
@@ -640,6 +630,37 @@ cdef class ScanOptions:
     cdef inline shared_ptr[CScanOptions] unwrap(self):
         return self.wrapped
 
+    @property
+    def schema(self):
+        if self.options.schema == nullptr:
+            return None
+        else:
+            return pyarrow_wrap_schema(self.options.schema)
+
+    @schema.setter
+    def schema(self, Schema value):
+        self.options.schema = pyarrow_unwrap_schema(value)
+
+
+# cdef class FileScanOptions(ScanOptions):
+#     pass
+
+
+# cdef class ParquetScanOptions(FileScanOptions):
+#     pass
+
+
+cdef class WriteOptions:
+    pass
+
+
+# cdef class FileWriteOptions(WriteOptions):
+#     pass
+
+
+# cdef class ParquetWriterOptions(FileWriteOptions):
+#     pass
+
 
 cdef class ScanContext:
 
@@ -649,11 +670,9 @@ cdef class ScanContext:
 
     def __init__(self, MemoryPool memory_pool=None):
         cdef shared_ptr[CScanContext] context = make_shared[CScanContext]()
-
-        if memory_pool is not None:
-            context.get().pool = memory_pool.pool
-
         self.init(context)
+        if memory_pool is not None:
+            self.memory_pool = memory_pool
 
     cdef init(self, shared_ptr[CScanContext]& sp):
         self.wrapped = sp
@@ -668,6 +687,19 @@ cdef class ScanContext:
     cdef inline shared_ptr[CScanContext] unwrap(self):
         return self.wrapped
 
+    @property
+    def memory_pool(self):
+        # TODO(kszucs): there's no function to box/wrap memory pool
+        cdef:
+            MemoryPool pool = MemoryPool.__new__(MemoryPool)
+        # this might be unsafe because of the raw pointers
+        pool.init(self.context.pool)
+        return pool
+
+    @memory_pool.setter
+    def memory_pool(self, MemoryPool pool):
+        self.context.pool = maybe_unbox_memory_pool(pool)
+
 
 cdef class ScanTask:
 

python/pyarrow/includes/libarrow_dataset.pxd

@@ -64,21 +64,6 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         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;
 
     ctypedef shared_ptr[CExpression] CExpressionPtr \
         "arrow::dataset::ExpressionPtr"
@@ -149,8 +134,8 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
 
     cdef cppclass CScanOptions "arrow::dataset::ScanOptions":
         CExpressionPtr filter
-        # shared_ptr[CExpressionEvaluator] evaluator
         shared_ptr[CSchema] schema
+        # shared_ptr[CExpressionEvaluator] evaluator
         # shared_ptr[CRecordBatchProjector] projector
 
         @staticmethod
@@ -195,6 +180,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     cdef cppclass CDataFragment "arrow::dataset::DataFragment":
         CResult[CScanTaskIterator] Scan(CScanContextPtr context)
         c_bool splittable()
+        c_string type()
         CScanOptionsPtr scan_options()
 
     ctypedef shared_ptr[CDataFragment] CDataFragmentPtr \
@@ -272,11 +258,10 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     ctypedef shared_ptr[CFileFormat] CFileFormatPtr \
         "arrow::dataset::FileFormatPtr"
 
-    cdef cppclass CFileBasedDataFragment \
-            "arrow::dataset::FileBasedDataFragment"(CDataFragment):
-        CFileBasedDataFragment(const CFileSource& source,
-                               CFileFormatPtr format,
-                               CScanOptionsPtr scan_options)
+    cdef cppclass CFileDataFragment "arrow::dataset::FileDataFragment"(
+            CDataFragment):
+        CFileDataFragment(const CFileSource& source, CFileFormatPtr format,
+                          CScanOptionsPtr scan_options)
         CStatus Scan(CScanContextPtr scan_context,
                      shared_ptr[CScanTaskIterator]* out)
         const CFileSource& source()
@@ -310,7 +295,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         pass
 
     cdef cppclass CParquetFragment "arrow::dataset::ParquetFragment"(
-            CFileBasedDataFragment):
+            CFileDataFragment):
         CParquetFragment(const CFileSource& source,
                          shared_ptr[CScanOptions] options)
         c_bool splittable()

python/pyarrow/tests/test_dataset.py

@@ -93,11 +93,23 @@ def dataset(simple_data_source, tree_data_source, schema):
 
 
 def test_scan_options():
-    pass
+    schema = pa.schema([
+        pa.field('a', pa.string()),
+        pa.field('b', pa.string())
+    ])
+    options = ds.ScanOptions(schema)
+    assert isinstance(options, ds.ScanOptions)
+    assert options.schema.equals(schema)
+
+    options = ds.ScanOptions()
+    assert options.schema is None
 
 
 def test_scan_context():
-    pass
+    context = ds.ScanContext()
+    assert isinstance(context.memory_pool, pa.MemoryPool)
+    context = ds.ScanContext(pa.default_memory_pool())
+    assert isinstance(context.memory_pool, pa.MemoryPool)
 
 
 def test_simple_data_fragment(record_batch):