Data fragments

cpp/src/arrow/dataset/file_base.h

@@ -158,8 +158,6 @@ 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/test_util.h

@@ -205,6 +205,8 @@ class DummyFragment : public FileDataFragment {
   DummyFragment(const FileSource& source, ScanOptionsPtr options)
       : FileDataFragment(source, std::make_shared<DummyFileFormat>(), options) {}
 
+  std::string type() const override { return "dummy_data_fragment"; }
+
   bool splittable() const override { return false; }
 };
 
@@ -261,6 +263,8 @@ class JSONRecordBatchFragment : public FileDataFragment {
       : FileDataFragment(source, std::make_shared<JSONRecordBatchFileFormat>(schema),
                          options) {}
 
+  std::string type() const override { return "json_batch_data_source"; }
+
   bool splittable() const override { return false; }
 };
 

python/pyarrow/_dataset.pyx

@@ -314,7 +314,7 @@ cdef class FileSource:
         return self.wrapped
 
     def equals(self, FileSource other):
-        return self.wrapped.get() == other.wrapped.get()
+        return deref(self.source) == deref(other.unwrap())
 
     def __eq__(self, other):
         try:
@@ -353,8 +353,6 @@ cdef class DataFragment:
         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:
@@ -415,11 +413,31 @@ cdef class SimpleDataFragment(DataFragment):
 
 
 cdef class FileDataFragment(DataFragment):
-    pass
+
+    cdef:
+        CFileDataFragment* file_fragment
+
+    cdef void init(self, const shared_ptr[CDataFragment]& sp):
+        DataFragment.init(self, sp)
+        self.file_fragment = <CFileDataFragment*> sp.get()
+
+    @property
+    def source(self):
+        cdef shared_ptr[CFileSource] source
+        source = make_shared[CFileSource](self.file_fragment.source())
+        return FileSource.wrap(source)
 
 
 cdef class ParquetDataFragment(FileDataFragment):
-    pass
+
+    def __init__(self, FileSource source not None,
+                 ParquetScanOptions scan_options not None):
+        cdef shared_ptr[CDataFragment] fragment
+        fragment.reset(new CParquetDataFragment(
+            deref(source.unwrap()),
+            scan_options.unwrap()
+        ))
+        self.init(fragment)
 
 
 cdef class DataSource:
@@ -642,24 +660,24 @@ cdef class ScanOptions:
         self.options.schema = pyarrow_unwrap_schema(value)
 
 
-# cdef class FileScanOptions(ScanOptions):
-#     pass
+cdef class FileScanOptions(ScanOptions):
+    pass
 
 
-# cdef class ParquetScanOptions(FileScanOptions):
-#     pass
+cdef class ParquetScanOptions(FileScanOptions):
+    pass
 
 
 cdef class WriteOptions:
     pass
 
 
-# cdef class FileWriteOptions(WriteOptions):
-#     pass
+cdef class FileWriteOptions(WriteOptions):
+    pass
 
 
-# cdef class ParquetWriterOptions(FileWriteOptions):
-#     pass
+cdef class ParquetWriterOptions(FileWriteOptions):
+    pass
 
 
 cdef class ScanContext:

python/pyarrow/dataset.py

@@ -33,7 +33,9 @@ from pyarrow._dataset import (  # noqa
     DataSource,
     Expression,
     FieldExpression,
+    FileDataFragment,
     FileFormat,
+    FileScanOptions,
     FileSource,
     FileSystemDataSource,
     FileSystemDataSourceDiscovery,
@@ -43,7 +45,9 @@ from pyarrow._dataset import (  # noqa
     IsValidExpression,
     NotExpression,
     OrExpression,
+    ParquetDataFragment,
     ParquetFileFormat,
+    ParquetScanOptions,
     PartitionScheme,
     ScalarExpression,
     ScanContext,

python/pyarrow/includes/libarrow_dataset.pxd

@@ -231,7 +231,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     cdef cppclass CFileSource "arrow::dataset::FileSource":
         CFileSource(c_string path, CFileSystem* filesystem,
                     CompressionType compression)
-        c_bool operator==(const CFileSource other)
+        c_bool operator==(const CFileSource& other) const
         CompressionType compression()
         c_string path()
         CFileSystem* filesystem()
@@ -268,6 +268,11 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CFileFormatPtr format()
         CScanOptionsPtr scan_options()
 
+    cdef cppclass CParquetDataFragment "arrow::dataset::ParquetFragment"(
+            CFileDataFragment):
+        CParquetDataFragment(const CFileSource& source,
+                             CScanOptionsPtr options)
+
     ctypedef unordered_map[c_string, CExpressionPtr] CPathPartitions \
         "arrow::dataset::PathPartitions"
 

python/pyarrow/tests/test_dataset.py

@@ -255,14 +255,24 @@ def test_scanner_builder(dataset):
 
 
 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'
+    source0 = ds.FileSource('/path/to/file.parquet', mockfs,
+                            compression=None)
+    source1 = ds.FileSource('/path/to/file.parquet', mockfs,
+                            compression=None)
+    source2 = ds.FileSource('/path/to/file.parquet.gz', mockfs,
+                            compression='gzip')
+    assert source1.path == '/path/to/file.parquet'
     assert source1.fs == mockfs
     assert source1.compression == 0  # None
-    assert source2.path == '/path/to/file.ext.gz'
+    assert source2.path == '/path/to/file.parquet.gz'
     assert source2.fs == mockfs
     assert source2.compression == 2  # 'gzip'
+    assert source1 != source2
+    assert source0 == source1
+
+    fragment = ds.ParquetDataFragment(source1, ds.ParquetScanOptions())
+    assert isinstance(fragment, ds.ParquetDataFragment)
+    assert fragment.source == source1
 
 
 def test_abstract_classes():