ARROW-6341: [Python] Implement low-level bindings for Dataset

LICENSE.txt

@@ -940,6 +940,33 @@ License: The MIT License (MIT)
 
 --------------------------------------------------------------------------------
 
+This project includes code from the cymove project:
+
+* python/pyarrow/includes/common.pxd includes code from the cymove project
+
+The MIT License (MIT)
+Copyright (c) 2019 Omer Ozarslan
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
+IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM,
+DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
+OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE
+OR OTHER DEALINGS IN THE SOFTWARE.
+
+--------------------------------------------------------------------------------
+
 This project include code from CMake.
 
 * cpp/cmake_modules/FindGTest.cmake is based on code from CMake.

ci/scripts/python_build.sh

@@ -32,6 +32,7 @@ export PYARROW_WITH_FLIGHT=${ARROW_FLIGHT:-OFF}
 export PYARROW_WITH_PLASMA=${ARROW_PLASMA:-OFF}
 export PYARROW_WITH_GANDIVA=${ARROW_GANDIVA:-OFF}
 export PYARROW_WITH_PARQUET=${ARROW_PARQUET:-OFF}
+export PYARROW_WITH_DATASET=${ARROW_DATASET:-OFF}
 
 export LD_LIBRARY_PATH=${ARROW_HOME}/lib:${LD_LIBRARY_PATH}
 

cpp/src/arrow/dataset/dataset.h

@@ -97,7 +97,8 @@ class ARROW_DS_EXPORT DataSource {
   /// May be null, which indicates no information is available.
   const ExpressionPtr& partition_expression() const { return partition_expression_; }
 
-  virtual std::string type() const = 0;
+  /// \brief The name identifying the kind of data source
+  virtual std::string type_name() const = 0;
 
   virtual ~DataSource() = default;
 
@@ -123,7 +124,7 @@ class ARROW_DS_EXPORT SimpleDataSource : public DataSource {
 
   DataFragmentIterator GetFragmentsImpl(ScanOptionsPtr options) override;
 
-  std::string type() const override { return "simple_data_source"; }
+  std::string type_name() const override { return "simple"; }
 
  private:
   DataFragmentVector fragments_;
@@ -136,7 +137,7 @@ class ARROW_DS_EXPORT TreeDataSource : public DataSource {
 
   DataFragmentIterator GetFragmentsImpl(ScanOptionsPtr options) override;
 
-  std::string type() const override { return "tree_data_source"; }
+  std::string type_name() const override { return "tree"; }
 
  private:
   DataSourceVector children_;
@@ -149,7 +150,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);
 

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,
@@ -130,7 +131,8 @@ class ARROW_DS_EXPORT FileFormat {
  public:
   virtual ~FileFormat() = default;
 
-  virtual std::string name() const = 0;
+  /// \brief The name identifying the kind of file format
+  virtual std::string type_name() const = 0;
 
   /// \brief Indicate if the FileSource is supported/readable by this format.
   virtual Result<bool> IsSupported(const FileSource& source) const = 0;
@@ -215,7 +217,7 @@ class ARROW_DS_EXPORT FileSystemDataSource : public DataSource {
                                     ExpressionVector partitions,
                                     ExpressionPtr source_partition, FileFormatPtr format);
 
-  std::string type() const override { return "filesystem_data_source"; }
+  std::string type_name() const override { return "filesystem"; }
 
   std::string ToString() const;
 

cpp/src/arrow/dataset/file_parquet.cc

@@ -50,7 +50,7 @@ class ParquetScanTask : public ScanTask {
         column_projection_(std::move(column_projection)),
         reader_(reader) {}
 
-  Result<RecordBatchIterator> Scan() {
+  Result<RecordBatchIterator> Execute() {
     // The construction of parquet's RecordBatchReader is deferred here to
     // control the memory usage of consumers who materialize all ScanTasks
     // before dispatching them, e.g. for scheduling purposes.

cpp/src/arrow/dataset/file_parquet.h

@@ -46,7 +46,7 @@ class ARROW_DS_EXPORT ParquetWriteOptions : public FileWriteOptions {
 /// \brief A FileFormat implementation that reads from Parquet files
 class ARROW_DS_EXPORT ParquetFileFormat : public FileFormat {
  public:
-  std::string name() const override { return "parquet"; }
+  std::string type_name() const override { return "parquet"; }
 
   Result<bool> IsSupported(const FileSource& source) const override;
 

cpp/src/arrow/dataset/file_parquet_test.cc

@@ -174,7 +174,7 @@ TEST_F(TestParquetFileFormat, ScanRecordBatchReader) {
 
   for (auto maybe_task : scan_task_it) {
     ASSERT_OK_AND_ASSIGN(auto task, std::move(maybe_task));
-    ASSERT_OK_AND_ASSIGN(auto rb_it, task->Scan());
+    ASSERT_OK_AND_ASSIGN(auto rb_it, task->Execute());
     for (auto maybe_batch : rb_it) {
       ASSERT_OK_AND_ASSIGN(auto batch, std::move(maybe_batch));
       row_count += batch->num_rows();
@@ -221,7 +221,7 @@ TEST_F(TestParquetFileFormat, ScanRecordBatchReaderProjected) {
 
   for (auto maybe_task : scan_task_it) {
     ASSERT_OK_AND_ASSIGN(auto task, std::move(maybe_task));
-    ASSERT_OK_AND_ASSIGN(auto rb_it, task->Scan());
+    ASSERT_OK_AND_ASSIGN(auto rb_it, task->Execute());
     for (auto maybe_batch : rb_it) {
       ASSERT_OK_AND_ASSIGN(auto batch, std::move(maybe_batch));
       row_count += batch->num_rows();
@@ -262,7 +262,7 @@ TEST_F(TestParquetFileFormat, ScanRecordBatchReaderProjectedMissingCols) {
 
   for (auto maybe_task : scan_task_it) {
     ASSERT_OK_AND_ASSIGN(auto task, std::move(maybe_task));
-    ASSERT_OK_AND_ASSIGN(auto rb_it, task->Scan());
+    ASSERT_OK_AND_ASSIGN(auto rb_it, task->Execute());
     for (auto maybe_batch : rb_it) {
       ASSERT_OK_AND_ASSIGN(auto batch, std::move(maybe_batch));
       row_count += batch->num_rows();
@@ -307,7 +307,7 @@ void CountRowsInScan(ScanTaskIterator& it, int64_t expected_rows,
 
   for (auto maybe_scan_task : it) {
     ASSERT_OK_AND_ASSIGN(auto scan_task, std::move(maybe_scan_task));
-    ASSERT_OK_AND_ASSIGN(auto rb_it, scan_task->Scan());
+    ASSERT_OK_AND_ASSIGN(auto rb_it, scan_task->Execute());
     for (auto maybe_record_batch : rb_it) {
       ASSERT_OK_AND_ASSIGN(auto record_batch, std::move(maybe_record_batch));
       actual_rows += record_batch->num_rows();
@@ -329,7 +329,7 @@ class TestParquetFileFormatPushDown : public TestParquetFileFormat {
     ASSERT_OK_AND_ASSIGN(auto it, fragment.Scan(ctx_));
     for (auto maybe_scan_task : it) {
       ASSERT_OK_AND_ASSIGN(auto scan_task, std::move(maybe_scan_task));
-      ASSERT_OK_AND_ASSIGN(auto rb_it, scan_task->Scan());
+      ASSERT_OK_AND_ASSIGN(auto rb_it, scan_task->Execute());
       for (auto maybe_record_batch : rb_it) {
         ASSERT_OK_AND_ASSIGN(auto record_batch, std::move(maybe_record_batch));
         actual_rows += record_batch->num_rows();

cpp/src/arrow/dataset/filter.h

@@ -184,6 +184,7 @@ class ARROW_DS_EXPORT Expression {
   /// returns a debug string representing this expression
   virtual std::string ToString() const = 0;
 
+  /// \brief Return the expression's type identifier
   ExpressionType::type type() const { return type_; }
 
   /// Copy this expression into a shared pointer.

cpp/src/arrow/dataset/partition.cc

@@ -54,17 +54,6 @@ Result<ExpressionPtr> PartitionScheme::Parse(const std::string& path) const {
   return and_(std::move(expressions));
 }
 
-class DefaultPartitionScheme : public PartitionScheme {
- public:
-  DefaultPartitionScheme() : PartitionScheme(::arrow::schema({})) {}
-
-  std::string name() const override { return "default_partition_scheme"; }
-
-  Result<ExpressionPtr> Parse(const std::string& segment, int i) const override {
-    return scalar(true);
-  }
-};
-
 PartitionSchemePtr PartitionScheme::Default() {
   return std::make_shared<DefaultPartitionScheme>();
 }

cpp/src/arrow/dataset/partition.h

@@ -60,7 +60,7 @@ class ARROW_DS_EXPORT PartitionScheme {
   virtual ~PartitionScheme() = default;
 
   /// \brief The name identifying the kind of partition scheme
-  virtual std::string name() const = 0;
+  virtual std::string type_name() const = 0;
 
   /// \brief Parse a path segment into a partition expression
   ///
@@ -108,6 +108,18 @@ class ARROW_DS_EXPORT PartitionSchemeDiscovery {
   std::shared_ptr<Schema> schema_;
 };
 
+/// \brief Subclass for representing the default, always true scheme.
+class DefaultPartitionScheme : public PartitionScheme {
+ public:
+  DefaultPartitionScheme() : PartitionScheme(::arrow::schema({})) {}
+
+  std::string type_name() const override { return "default"; }
+
+  Result<ExpressionPtr> Parse(const std::string& segment, int i) const override {
+    return scalar(true);
+  }
+};
+
 /// \brief Subclass for looking up partition information from a dictionary
 /// mapping segments to expressions provided on construction.
 class ARROW_DS_EXPORT SegmentDictionaryPartitionScheme : public PartitionScheme {
@@ -117,7 +129,7 @@ class ARROW_DS_EXPORT SegmentDictionaryPartitionScheme : public PartitionScheme
       std::vector<std::unordered_map<std::string, ExpressionPtr>> dictionaries)
       : PartitionScheme(std::move(schema)), dictionaries_(std::move(dictionaries)) {}
 
-  std::string name() const override { return "segment_dictionary_partition_scheme"; }
+  std::string type_name() const override { return "segment_dictionary"; }
 
   /// Return dictionaries_[i][segment] or scalar(true)
   Result<ExpressionPtr> Parse(const std::string& segment, int i) const override;
@@ -160,7 +172,7 @@ class ARROW_DS_EXPORT SchemaPartitionScheme : public PartitionKeysScheme {
   explicit SchemaPartitionScheme(std::shared_ptr<Schema> schema)
       : PartitionKeysScheme(std::move(schema)) {}
 
-  std::string name() const override { return "schema_partition_scheme"; }
+  std::string type_name() const override { return "schema"; }
 
   util::optional<Key> ParseKey(const std::string& segment, int i) const override;
 
@@ -181,7 +193,7 @@ class ARROW_DS_EXPORT HivePartitionScheme : public PartitionKeysScheme {
   explicit HivePartitionScheme(std::shared_ptr<Schema> schema)
       : PartitionKeysScheme(std::move(schema)) {}
 
-  std::string name() const override { return "hive_partition_scheme"; }
+  std::string type_name() const override { return "hive"; }
 
   util::optional<Key> ParseKey(const std::string& segment, int i) const override {
     return ParseKey(segment);
@@ -198,12 +210,12 @@ class ARROW_DS_EXPORT FunctionPartitionScheme : public PartitionScheme {
   explicit FunctionPartitionScheme(
       std::shared_ptr<Schema> schema,
       std::function<Result<ExpressionPtr>(const std::string&, int)> impl,
-      std::string name = "function_partition_scheme")
+      std::string name = "function")
       : PartitionScheme(std::move(schema)),
         impl_(std::move(impl)),
         name_(std::move(name)) {}
 
-  std::string name() const override { return name_; }
+  std::string type_name() const override { return name_; }
 
   Result<ExpressionPtr> Parse(const std::string& segment, int i) const override {
     return impl_(segment, i);

cpp/src/arrow/dataset/partition_test.cc

@@ -229,7 +229,7 @@ class RangePartitionScheme : public HivePartitionScheme {
  public:
   using HivePartitionScheme::HivePartitionScheme;
 
-  std::string name() const override { return "range_partition_scheme"; }
+  std::string type_name() const override { return "range"; }
 
   Result<ExpressionPtr> Parse(const std::string& segment, int i) const override {
     ExpressionVector ranges;

cpp/src/arrow/dataset/scanner.cc

@@ -46,7 +46,7 @@ ScanOptionsPtr ScanOptions::ReplaceSchema(std::shared_ptr<Schema> schema) const
   return copy;
 }
 
-Result<RecordBatchIterator> SimpleScanTask::Scan() {
+Result<RecordBatchIterator> SimpleScanTask::Execute() {
   return MakeVectorIterator(record_batches_);
 }
 
@@ -169,7 +169,7 @@ struct TableAggregator {
 
 struct ScanTaskPromise {
   Status operator()() {
-    ARROW_ASSIGN_OR_RAISE(auto it, task->Scan());
+    ARROW_ASSIGN_OR_RAISE(auto it, task->Execute());
     for (auto maybe_batch : it) {
       ARROW_ASSIGN_OR_RAISE(auto batch, std::move(maybe_batch));
       aggregator.Append(std::move(batch));

cpp/src/arrow/dataset/scanner.h

@@ -83,9 +83,9 @@ class ARROW_DS_EXPORT ScanOptions {
 class ARROW_DS_EXPORT ScanTask {
  public:
   /// \brief Iterate through sequence of materialized record batches
-  /// resulting from the Scan. Execution semantics encapsulated in the
+  /// resulting from the Scan. Execution semantics are encapsulated in the
   /// particular ScanTask implementation
-  virtual Result<RecordBatchIterator> Scan() = 0;
+  virtual Result<RecordBatchIterator> Execute() = 0;
 
   virtual ~ScanTask() = default;
 
@@ -108,7 +108,7 @@ class ARROW_DS_EXPORT SimpleScanTask : public ScanTask {
       : ScanTask(std::move(options), std::move(context)),
         record_batches_(std::move(record_batches)) {}
 
-  Result<RecordBatchIterator> Scan() override;
+  Result<RecordBatchIterator> Execute() override;
 
  protected:
   std::vector<std::shared_ptr<RecordBatch>> record_batches_;
@@ -165,7 +165,7 @@ class ARROW_DS_EXPORT ScannerBuilder {
 
   /// \brief Set the subset of columns to materialize.
   ///
-  /// This subset wil be passed down to DataSources and corresponding DataFragments.
+  /// This subset will be passed down to DataSources and corresponding DataFragments.
   /// The goal is to avoid loading/copying/deserializing columns that will
   /// not be required further down the compute chain.
   ///
@@ -181,7 +181,6 @@ class ARROW_DS_EXPORT ScannerBuilder {
   /// The predicate will be passed down to DataSources and corresponding
   /// DataFragments to exploit predicate pushdown if possible using
   /// partition information or DataFragment internal metadata, e.g. Parquet statistics.
-  /// statistics.
   ///
   /// \param[in] filter expression to filter rows with.
   ///

cpp/src/arrow/dataset/scanner_internal.h

@@ -53,8 +53,8 @@ class FilterAndProjectScanTask : public ScanTask {
   explicit FilterAndProjectScanTask(ScanTaskPtr task)
       : ScanTask(task->options(), task->context()), task_(std::move(task)) {}
 
-  Result<RecordBatchIterator> Scan() override {
-    ARROW_ASSIGN_OR_RAISE(auto it, task_->Scan());
+  Result<RecordBatchIterator> Execute() override {
+    ARROW_ASSIGN_OR_RAISE(auto it, task_->Execute());
     auto filter_it = FilterRecordBatch(std::move(it), *options_->evaluator,
                                        *options_->filter, context_->pool);
     return ProjectRecordBatch(std::move(filter_it), &task_->options()->projector,

cpp/src/arrow/dataset/test_util.h

@@ -87,7 +87,7 @@ class DatasetFixtureMixin : public ::testing::Test {
   /// record batches yielded by the data fragment.
   void AssertScanTaskEquals(RecordBatchReader* expected, ScanTask* task,
                             bool ensure_drained = true) {
-    ASSERT_OK_AND_ASSIGN(auto it, task->Scan());
+    ASSERT_OK_AND_ASSIGN(auto it, task->Execute());
     ARROW_EXPECT_OK(it.Visit([expected](std::shared_ptr<RecordBatch> rhs) -> Status {
       std::shared_ptr<RecordBatch> lhs;
       RETURN_NOT_OK(expected->ReadNext(&lhs));
@@ -179,7 +179,7 @@ class DummyFileFormat : public FileFormat {
   explicit DummyFileFormat(std::shared_ptr<Schema> schema = NULLPTR)
       : schema_(std::move(schema)) {}
 
-  std::string name() const override { return "dummy"; }
+  std::string type_name() const override { return "dummy"; }
 
   Result<bool> IsSupported(const FileSource& source) const override { return true; }
 
@@ -223,7 +223,7 @@ class JSONRecordBatchFileFormat : public FileFormat {
   explicit JSONRecordBatchFileFormat(SchemaResolver resolver)
       : resolver_(std::move(resolver)) {}
 
-  std::string name() const override { return "json_record_batch"; }
+  std::string type_name() const override { return "json_record_batch"; }
 
   /// \brief Return true if the given file extension
   Result<bool> IsSupported(const FileSource& source) const override { return true; }

cpp/src/arrow/filesystem/filesystem.h

@@ -155,6 +155,8 @@ class ARROW_EXPORT FileSystem {
  public:
   virtual ~FileSystem();
 
+  virtual std::string type_name() const = 0;
+
   /// Get statistics for the given target.
   ///
   /// Any symlink is automatically dereferenced, recursively.
@@ -243,6 +245,8 @@ class ARROW_EXPORT SubTreeFileSystem : public FileSystem {
                              std::shared_ptr<FileSystem> base_fs);
   ~SubTreeFileSystem() override;
 
+  std::string type_name() const override { return "subtree"; }
+
   /// \cond FALSE
   using FileSystem::GetTargetStats;
   /// \endcond
@@ -289,6 +293,8 @@ class ARROW_EXPORT SlowFileSystem : public FileSystem {
   SlowFileSystem(std::shared_ptr<FileSystem> base_fs, double average_latency,
                  int32_t seed);
 
+  std::string type_name() const override { return "slow"; }
+
   using FileSystem::GetTargetStats;
   Result<FileStats> GetTargetStats(const std::string& path) override;
   Result<std::vector<FileStats>> GetTargetStats(const Selector& select) override;