text: // 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.
  
  #pragma once
  
  #include <algorithm>
  #include <ciso646>
  #include <functional>
  #include <memory>
  #include <ostream>
  #include <string>
  #include <unordered_map>
  #include <unordered_set>
  #include <utility>
  #include <vector>
  
  #include <gmock/gmock.h>
  #include <gtest/gtest.h>
  
  #include "arrow/array.h"
  #include "arrow/compute/exec/expression.h"
  #include "arrow/dataset/dataset_internal.h"
  #include "arrow/dataset/discovery.h"
  #include "arrow/dataset/file_base.h"
  #include "arrow/filesystem/localfs.h"
  #include "arrow/filesystem/mockfs.h"
  #include "arrow/filesystem/path_util.h"
  #include "arrow/filesystem/test_util.h"
  #include "arrow/record_batch.h"
  #include "arrow/table.h"
  #include "arrow/testing/future_util.h"
  #include "arrow/testing/generator.h"
  #include "arrow/testing/gtest_util.h"
  #include "arrow/testing/matchers.h"
  #include "arrow/testing/random.h"
  #include "arrow/util/async_generator.h"
  #include "arrow/util/io_util.h"
  #include "arrow/util/iterator.h"
  #include "arrow/util/logging.h"
  #include "arrow/util/make_unique.h"
  #include "arrow/util/thread_pool.h"
  
  namespace arrow {
  
  using internal::checked_cast;
  using internal::checked_pointer_cast;
  using internal::TemporaryDir;
  
  namespace dataset {
  
  using compute::call;
  using compute::field_ref;
  using compute::literal;
  
  using compute::and_;
  using compute::equal;
  using compute::greater;
  using compute::greater_equal;
  using compute::is_null;
  using compute::is_valid;
  using compute::less;
  using compute::less_equal;
  using compute::not_;
  using compute::not_equal;
  using compute::or_;
  using compute::project;
  
  using fs::internal::GetAbstractPathExtension;
  
  /// \brief Assert a dataset produces data with the schema
  void AssertDatasetHasSchema(std::shared_ptr<Dataset> ds, std::shared_ptr<Schema> schema) {
    ASSERT_OK_AND_ASSIGN(auto scanner_builder, ds->NewScan());
    ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder->Finish());
    ASSERT_OK_AND_ASSIGN(auto table, scanner->ToTable());
    ASSERT_EQ(*table->schema(), *schema);
  }
  
  class FileSourceFixtureMixin : public ::testing::Test {
   public:
    std::unique_ptr<FileSource> GetSource(std::shared_ptr<Buffer> buffer) {
      return ::arrow::internal::make_unique<FileSource>(std::move(buffer));
    }
  };
  
  template <typename Gen>
  class GeneratedRecordBatch : public RecordBatchReader {
   public:
    GeneratedRecordBatch(std::shared_ptr<Schema> schema, Gen gen)
        : schema_(std::move(schema)), gen_(gen) {}
  
    std::shared_ptr<Schema> schema() const override { return schema_; }
  
    Status ReadNext(std::shared_ptr<RecordBatch>* batch) override { return gen_(batch); }
  
   private:
    std::shared_ptr<Schema> schema_;
    Gen gen_;
  };
  
  template <typename Gen>
  std::unique_ptr<GeneratedRecordBatch<Gen>> MakeGeneratedRecordBatch(
      std::shared_ptr<Schema> schema, Gen&& gen) {
    return ::arrow::internal::make_unique<GeneratedRecordBatch<Gen>>(
        schema, std::forward<Gen>(gen));
  }
  
  std::unique_ptr<RecordBatchReader> MakeGeneratedRecordBatch(
      std::shared_ptr<Schema> schema, int64_t batch_size, int64_t batch_repetitions) {
    auto batch = random::GenerateBatch(schema->fields(), batch_size, /*seed=*/0);
    int64_t i = 0;
    return MakeGeneratedRecordBatch(
        schema, [batch, i, batch_repetitions](std::shared_ptr<RecordBatch>* out) mutable {
          *out = i++ < batch_repetitions ? batch : nullptr;
          return Status::OK();
        });
  }
  
  void EnsureRecordBatchReaderDrained(RecordBatchReader* reader) {
    ASSERT_OK_AND_ASSIGN(auto batch, reader->Next());
    EXPECT_EQ(batch, nullptr);
  }
  
  class DatasetFixtureMixin : public ::testing::Test {
   public:
    /// \brief Ensure that record batches found in reader are equals to the
    /// record batches yielded by the data fragment.
    void AssertScanTaskEquals(RecordBatchReader* expected, RecordBatchGenerator batch_gen,
                              bool ensure_drained = true) {
      ASSERT_FINISHES_OK(VisitAsyncGenerator(
          batch_gen, [expected](std::shared_ptr<RecordBatch> rhs) -> Status {
            std::shared_ptr<RecordBatch> lhs;
            RETURN_NOT_OK(expected->ReadNext(&lhs));
            EXPECT_NE(lhs, nullptr);
            AssertBatchesEqual(*lhs, *rhs);
            return Status::OK();
          }));
  
      if (ensure_drained) {
        EnsureRecordBatchReaderDrained(expected);
      }
    }
  
    /// \brief Assert the value of the next batch yielded by the reader
    void AssertBatchEquals(RecordBatchReader* expected, const RecordBatch& batch) {
      std::shared_ptr<RecordBatch> lhs;
      ASSERT_OK(expected->ReadNext(&lhs));
      EXPECT_NE(lhs, nullptr);
      AssertBatchesEqual(*lhs, batch);
    }
  
    /// \brief Ensure that record batches found in reader are equals to the
    /// record batches yielded by the data fragment.
    void AssertFragmentEquals(RecordBatchReader* expected, Fragment* fragment,
                              bool ensure_drained = true) {
      ASSERT_OK_AND_ASSIGN(auto batch_gen, fragment->ScanBatchesAsync(options_));
      AssertScanTaskEquals(expected, batch_gen);
  
      if (ensure_drained) {
        EnsureRecordBatchReaderDrained(expected);
      }
    }
  
    /// \brief Ensure that record batches found in reader are equals to the
    /// record batches yielded by the data fragments of a dataset.
    void AssertDatasetFragmentsEqual(RecordBatchReader* expected, Dataset* dataset,
                                     bool ensure_drained = true) {
      ASSERT_OK_AND_ASSIGN(auto predicate, options_->filter.Bind(*dataset->schema()));
      ASSERT_OK_AND_ASSIGN(auto it, dataset->GetFragments(predicate));
  
      ARROW_EXPECT_OK(it.Visit([&](std::shared_ptr<Fragment> fragment) -> Status {
        AssertFragmentEquals(expected, fragment.get(), false);
        return Status::OK();
      }));
  
      if (ensure_drained) {
        EnsureRecordBatchReaderDrained(expected);
      }
    }
  
    /// \brief Ensure that record batches found in reader are equals to the
    /// record batches yielded by a scanner.
    void AssertScannerEquals(RecordBatchReader* expected, Scanner* scanner,
                             bool ensure_drained = true) {
      ASSERT_OK_AND_ASSIGN(auto it, scanner->ScanBatches());
  
      ARROW_EXPECT_OK(it.Visit([&](TaggedRecordBatch batch) -> Status {
        std::shared_ptr<RecordBatch> lhs;
        RETURN_NOT_OK(expected->ReadNext(&lhs));
        EXPECT_NE(lhs, nullptr);
        AssertBatchesEqual(*lhs, *batch.record_batch);
        return Status::OK();
      }));
  
      if (ensure_drained) {
        EnsureRecordBatchReaderDrained(expected);
      }
    }
  
    /// \brief Ensure that record batches found in reader are equals to the
    /// record batches yielded by a scanner.
    void AssertScanBatchesEquals(RecordBatchReader* expected, Scanner* scanner,
                                 bool ensure_drained = true) {
      ASSERT_OK_AND_ASSIGN(auto it, scanner->ScanBatches());
  
      ARROW_EXPECT_OK(it.Visit([&](TaggedRecordBatch batch) -> Status {
        AssertBatchEquals(expected, *batch.record_batch);
        return Status::OK();
      }));
  
      if (ensure_drained) {
        EnsureRecordBatchReaderDrained(expected);
      }
    }
  
    /// \brief Ensure that record batches found in reader are equals to the
    /// record batches yielded by a scanner.
    void AssertScanBatchesUnorderedEquals(RecordBatchReader* expected, Scanner* scanner,
                                          int expected_batches_per_fragment,
                                          bool ensure_drained = true) {
      ASSERT_OK_AND_ASSIGN(auto it, scanner->ScanBatchesUnordered());
  
      // ToVector does not work since EnumeratedRecordBatch is not comparable
      std::vector<EnumeratedRecordBatch> batches;
      for (;;) {
        ASSERT_OK_AND_ASSIGN(auto batch, it.Next());
        if (IsIterationEnd(batch)) break;
        batches.push_back(std::move(batch));
      }
      std::sort(batches.begin(), batches.end(),
                [](const EnumeratedRecordBatch& left,
                   const EnumeratedRecordBatch& right) -> bool {
                  if (left.fragment.index < right.fragment.index) {
                    return true;
                  }
                  if (left.fragment.index > right.fragment.index) {
                    return false;
                  }
                  return left.record_batch.index < right.record_batch.index;
                });
  
      int fragment_counter = 0;
      bool saw_last_fragment = false;
      int batch_counter = 0;
  
      for (const auto& batch : batches) {
        if (batch_counter == 0) {
          EXPECT_FALSE(saw_last_fragment);
        }
        EXPECT_EQ(batch_counter++, batch.record_batch.index);
        auto last_batch = batch_counter == expected_batches_per_fragment;
        EXPECT_EQ(last_batch, batch.record_batch.last);
        EXPECT_EQ(fragment_counter, batch.fragment.index);
        if (last_batch) {
          fragment_counter++;
          batch_counter = 0;
        }
        saw_last_fragment = batch.fragment.last;
        AssertBatchEquals(expected, *batch.record_batch.value);
      }
  
      if (ensure_drained) {
        EnsureRecordBatchReaderDrained(expected);
      }
    }
  
    /// \brief Ensure that record batches found in reader are equals to the
    /// record batches yielded by a dataset.
    void AssertDatasetEquals(RecordBatchReader* expected, Dataset* dataset,
                             bool ensure_drained = true) {
      ASSERT_OK_AND_ASSIGN(auto builder, dataset->NewScan());
      ASSERT_OK_AND_ASSIGN(auto scanner, builder->Finish());
      AssertScannerEquals(expected, scanner.get());
  
      if (ensure_drained) {
        EnsureRecordBatchReaderDrained(expected);
      }
    }
  
   protected:
    void SetSchema(std::vector<std::shared_ptr<Field>> fields) {
      schema_ = schema(std::move(fields));
      options_ = std::make_shared<ScanOptions>();
      options_->dataset_schema = schema_;
      ASSERT_OK_AND_ASSIGN(auto projection,
                           ProjectionDescr::FromNames(schema_->field_names(), *schema_));
      SetProjection(options_.get(), std::move(projection));
      SetFilter(literal(true));
    }
  
    void SetFilter(compute::Expression filter) {
      ASSERT_OK_AND_ASSIGN(options_->filter, filter.Bind(*schema_));
    }
  
    void SetProjectedColumns(std::vector<std::string> column_names) {
      ASSERT_OK_AND_ASSIGN(
          auto projection,
          ProjectionDescr::FromNames(std::move(column_names), *options_->dataset_schema));
      SetProjection(options_.get(), std::move(projection));
    }
  
    std::shared_ptr<Schema> schema_;
    std::shared_ptr<ScanOptions> options_;
  };
  
  template <typename P>
  class DatasetFixtureMixinWithParam : public DatasetFixtureMixin,
                                       public ::testing::WithParamInterface<P> {};
  
  struct TestFormatParams {
    bool use_threads;
    int num_batches;
    int items_per_batch;
  
    int64_t expected_rows() const { return num_batches * items_per_batch; }
  
    std::string ToString() const {
      // GTest requires this to be alphanumeric
      std::stringstream ss;
      ss << (use_threads ? "Threaded" : "Serial") << num_batches << "b" << items_per_batch
         << "r";
      return ss.str();
    }
  
    static std::string ToTestNameString(
        const ::testing::TestParamInfo<TestFormatParams>& info) {
      return std::to_string(info.index) + info.param.ToString();
    }
  
    static std::vector<TestFormatParams> Values() {
      std::vector<TestFormatParams> values;
      for (const bool use_threads : std::vector<bool>{true, false}) {
        values.push_back(TestFormatParams{use_threads, 16, 1024});
      }
      return values;
    }
  };
  
  std::ostream& operator<<(std::ostream& out, const TestFormatParams& params) {
    out << params.ToString();
    return out;
  }
  
  class FileFormatWriterMixin {
    virtual std::shared_ptr<Buffer> Write(RecordBatchReader* reader) = 0;
    virtual std::shared_ptr<Buffer> Write(const Table& table) = 0;
  };
  
  /// FormatHelper should be a class with these static methods:
  /// std::shared_ptr<Buffer> Write(RecordBatchReader* reader);
  /// std::shared_ptr<FileFormat> MakeFormat();
  template <typename FormatHelper>
  class FileFormatFixtureMixin : public ::testing::Test {
   public:
    constexpr static int64_t kBatchSize = 1UL << 12;
    constexpr static int64_t kBatchRepetitions = 1 << 5;
  
    FileFormatFixtureMixin()
        : format_(FormatHelper::MakeFormat()), opts_(std::make_shared<ScanOptions>()) {}
  
    int64_t expected_batches() const { return kBatchRepetitions; }
    int64_t expected_rows() const { return kBatchSize * kBatchRepetitions; }
  
    std::shared_ptr<FileFragment> MakeFragment(const FileSource& source) {
      EXPECT_OK_AND_ASSIGN(auto fragment, format_->MakeFragment(source));
      return fragment;
    }
  
    std::shared_ptr<FileFragment> MakeFragment(const FileSource& source,
                                               compute::Expression partition_expression) {
      EXPECT_OK_AND_ASSIGN(auto fragment,
                           format_->MakeFragment(source, partition_expression));
      return fragment;
    }
  
    std::shared_ptr<FileSource> GetFileSource(RecordBatchReader* reader) {
      EXPECT_OK_AND_ASSIGN(auto buffer, FormatHelper::Write(reader));
      return std::make_shared<FileSource>(std::move(buffer));
    }
  
    virtual std::shared_ptr<RecordBatchReader> GetRecordBatchReader(
        std::shared_ptr<Schema> schema) {
      return MakeGeneratedRecordBatch(schema, kBatchSize, kBatchRepetitions);
    }
  
    Result<std::shared_ptr<io::BufferOutputStream>> GetFileSink() {
      ARROW_ASSIGN_OR_RAISE(std::shared_ptr<ResizableBuffer> buffer,
                            AllocateResizableBuffer(0));
      return std::make_shared<io::BufferOutputStream>(buffer);
    }
  
    void SetSchema(std::vector<std::shared_ptr<Field>> fields) {
      opts_->dataset_schema = schema(std::move(fields));
      ASSERT_OK_AND_ASSIGN(auto projection,
                           ProjectionDescr::Default(*opts_->dataset_schema));
      SetProjection(opts_.get(), std::move(projection));
    }
  
    void SetFilter(compute::Expression filter) {
      ASSERT_OK_AND_ASSIGN(opts_->filter, filter.Bind(*opts_->dataset_schema));
    }
  
    void Project(std::vector<std::string> names) {
      ASSERT_OK_AND_ASSIGN(auto projection, ProjectionDescr::FromNames(
                                                std::move(names), *opts_->dataset_schema));
      SetProjection(opts_.get(), std::move(projection));
    }
  
    void ProjectNested(std::vector<std::string> names) {
      std::vector<compute::Expression> exprs;
      for (const auto& name : names) {
        ASSERT_OK_AND_ASSIGN(auto ref, FieldRef::FromDotPath(name));
        exprs.push_back(field_ref(ref));
      }
      ASSERT_OK_AND_ASSIGN(
          auto descr, ProjectionDescr::FromExpressions(std::move(exprs), std::move(names),
                                                       *opts_->dataset_schema));
      SetProjection(opts_.get(), std::move(descr));
    }
  
    // Shared test cases
    void AssertInspectFailure(const std::string& contents, StatusCode code,
                              const std::string& format_name) {
      SCOPED_TRACE("Format: " + format_name + " File contents: " + contents);
      constexpr auto file_name = "herp/derp";
      auto make_error_message = [&](const std::string& filename) {
        return "Could not open " + format_name + " input source '" + filename + "':";
      };
      const auto buf = std::make_shared<Buffer>(contents);
      Status status;
  
      status = format_->Inspect(FileSource(buf)).status();
      EXPECT_EQ(code, status.code());
      EXPECT_THAT(status.ToString(), ::testing::HasSubstr(make_error_message("<Buffer>")));
  
      ASSERT_OK_AND_EQ(false, format_->IsSupported(FileSource(buf)));
  
      ASSERT_OK_AND_ASSIGN(
          auto fs, fs::internal::MockFileSystem::Make(fs::kNoTime, {fs::File(file_name)}));
      status = format_->Inspect({file_name, fs}).status();
      EXPECT_EQ(code, status.code());
      EXPECT_THAT(status.ToString(), testing::HasSubstr(make_error_message("herp/derp")));
  
      fs::FileSelector s;
      s.base_dir = "/";
      s.recursive = true;
      FileSystemFactoryOptions options;
      ASSERT_OK_AND_ASSIGN(auto factory,
                           FileSystemDatasetFactory::Make(fs, s, format_, options));
      status = factory->Finish().status();
      EXPECT_EQ(code, status.code());
      EXPECT_THAT(
          status.ToString(),
          ::testing::AllOf(
              ::testing::HasSubstr(make_error_message("/herp/derp")),
              ::testing::HasSubstr(
                  "Error creating dataset. Could not read schema from '/herp/derp':"),
              ::testing::HasSubstr("Is this a '" + format_->type_name() + "' file?")));
    }
  
    void TestInspectFailureWithRelevantError(StatusCode code,
                                             const std::string& format_name) {
      const std::vector<std::string> file_contents{"", "PAR0", "ASDFPAR1", "ARROW1"};
      for (const auto& contents : file_contents) {
        AssertInspectFailure(contents, code, format_name);
      }
    }
  
    void TestInspect() {
      auto reader = GetRecordBatchReader(schema({field("f64", float64())}));
      auto source = GetFileSource(reader.get());
  
      ASSERT_OK_AND_ASSIGN(auto actual, format_->Inspect(*source.get()));
      AssertSchemaEqual(*actual, *reader->schema(), /*check_metadata=*/false);
    }
    void TestIsSupported() {
      auto reader = GetRecordBatchReader(schema({field("f64", float64())}));
      auto source = GetFileSource(reader.get());
  
      bool supported = false;
  
      std::shared_ptr<Buffer> buf = std::make_shared<Buffer>(util::string_view(""));
      ASSERT_OK_AND_ASSIGN(supported, format_->IsSupported(FileSource(buf)));
      ASSERT_EQ(supported, false);
  
      buf = std::make_shared<Buffer>(util::string_view("corrupted"));
      ASSERT_OK_AND_ASSIGN(supported, format_->IsSupported(FileSource(buf)));
      ASSERT_EQ(supported, false);
  
      ASSERT_OK_AND_ASSIGN(supported, format_->IsSupported(*source));
      EXPECT_EQ(supported, true);
    }
    std::shared_ptr<Buffer> WriteToBuffer(
        std::shared_ptr<Schema> schema,
        std::shared_ptr<FileWriteOptions> options = nullptr) {
      auto format = format_;
      SetSchema(schema->fields());
      EXPECT_OK_AND_ASSIGN(auto sink, GetFileSink());
      if (!options) options = format->DefaultWriteOptions();
  
      EXPECT_OK_AND_ASSIGN(auto fs, fs::internal::MockFileSystem::Make(fs::kNoTime, {}));
      EXPECT_OK_AND_ASSIGN(auto writer,
                           format->MakeWriter(sink, schema, options, {fs, "<buffer>"}));
      ARROW_EXPECT_OK(writer->Write(GetRecordBatchReader(schema).get()));
      auto fut = writer->Finish();
      EXPECT_FINISHES(fut);
      ARROW_EXPECT_OK(fut.status());
      EXPECT_OK_AND_ASSIGN(auto written, sink->Finish());
      return written;
    }
    void TestWrite() {
      auto reader = this->GetRecordBatchReader(schema({field("f64", float64())}));
      auto source = this->GetFileSource(reader.get());
      auto written = this->WriteToBuffer(reader->schema());
      AssertBufferEqual(*written, *source->buffer());
    }
    void TestCountRows() {
      auto options = std::make_shared<ScanOptions>();
      auto reader = this->GetRecordBatchReader(schema({field("f64", float64())}));
      auto full_schema = schema({field("f64", float64()), field("part", int64())});
      auto source = this->GetFileSource(reader.get());
  
      auto fragment = this->MakeFragment(*source);
      ASSERT_FINISHES_OK_AND_EQ(util::make_optional<int64_t>(expected_rows()),
                                fragment->CountRows(literal(true), options));
  
      fragment = this->MakeFragment(*source, equal(field_ref("part"), literal(2)));
      ASSERT_FINISHES_OK_AND_EQ(util::make_optional<int64_t>(expected_rows()),
                                fragment->CountRows(literal(true), options));
  
      auto predicate = equal(field_ref("part"), literal(1));
      ASSERT_OK_AND_ASSIGN(predicate, predicate.Bind(*full_schema));
      ASSERT_FINISHES_OK_AND_EQ(util::make_optional<int64_t>(0),
                                fragment->CountRows(predicate, options));
  
      predicate = equal(field_ref("part"), literal(2));
      ASSERT_OK_AND_ASSIGN(predicate, predicate.Bind(*full_schema));
      ASSERT_FINISHES_OK_AND_EQ(util::make_optional<int64_t>(expected_rows()),
                                fragment->CountRows(predicate, options));
  
      predicate = equal(call("add", {field_ref("f64"), literal(3)}), literal(2));
      ASSERT_OK_AND_ASSIGN(predicate, predicate.Bind(*full_schema));
      ASSERT_FINISHES_OK_AND_EQ(util::nullopt, fragment->CountRows(predicate, options));
    }
  
   protected:
    std::shared_ptr<typename FormatHelper::FormatType> format_;
    std::shared_ptr<ScanOptions> opts_;
  };
  
  template <typename FormatHelper>
  class FileFormatScanMixin : public FileFormatFixtureMixin<FormatHelper>,
                              public ::testing::WithParamInterface<TestFormatParams> {
   public:
    int64_t expected_batches() const { return GetParam().num_batches; }
    int64_t expected_rows() const { return GetParam().expected_rows(); }
  
    std::shared_ptr<RecordBatchReader> GetRecordBatchReader(
        std::shared_ptr<Schema> schema) override {
      return MakeGeneratedRecordBatch(schema, GetParam().items_per_batch,
                                      GetParam().num_batches);
    }
  
    // Scan the fragment through the scanner.
    RecordBatchIterator Batches(std::shared_ptr<Fragment> fragment) {
      auto dataset = std::make_shared<FragmentDataset>(opts_->dataset_schema,
                                                       FragmentVector{fragment});
      ScannerBuilder builder(dataset, opts_);
      ARROW_EXPECT_OK(builder.UseThreads(GetParam().use_threads));
      EXPECT_OK_AND_ASSIGN(auto scanner, builder.Finish());
      EXPECT_OK_AND_ASSIGN(auto batch_it, scanner->ScanBatches());
      return MakeMapIterator([](TaggedRecordBatch tagged) { return tagged.record_batch; },
                             std::move(batch_it));
    }
  
    // Scan the fragment directly, without using the scanner.
    RecordBatchIterator PhysicalBatches(std::shared_ptr<Fragment> fragment) {
      opts_->use_threads = GetParam().use_threads;
      EXPECT_OK_AND_ASSIGN(auto batch_gen, fragment->ScanBatchesAsync(opts_));
      auto batch_it = MakeGeneratorIterator(std::move(batch_gen));
      return batch_it;
    }
  
    // Shared test cases
    void TestScan() {
      auto reader = GetRecordBatchReader(schema({field("f64", float64())}));
      auto source = this->GetFileSource(reader.get());
  
      this->SetSchema(reader->schema()->fields());
      auto fragment = this->MakeFragment(*source);
  
      int64_t row_count = 0;
      for (auto maybe_batch : Batches(fragment)) {
        ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
        row_count += batch->num_rows();
      }
      ASSERT_EQ(row_count, GetParam().expected_rows());
    }
    // Ensure batch_size is respected
    void TestScanBatchSize() {
      constexpr int kBatchSize = 17;
      auto reader = GetRecordBatchReader(schema({field("f64", float64())}));
      auto source = this->GetFileSource(reader.get());
  
      this->SetSchema(reader->schema()->fields());
      auto fragment = this->MakeFragment(*source);
  
      int64_t row_count = 0;
      opts_->batch_size = kBatchSize;
      for (auto maybe_batch : Batches(fragment)) {
        ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
        ASSERT_LE(batch->num_rows(), kBatchSize);
        row_count += batch->num_rows();
      }
      ASSERT_EQ(row_count, GetParam().expected_rows());
    }
    // Ensure file formats only return columns needed to fulfill filter/projection
    void TestScanProjected() {
      auto f32 = field("f32", float32());
      auto f64 = field("f64", float64());
      auto i32 = field("i32", int32());
      auto i64 = field("i64", int64());
      this->SetSchema({f64, i64, f32, i32});
      this->Project({"f64"});
      this->SetFilter(equal(field_ref("i32"), literal(0)));
  
      // NB: projection is applied by the scanner; FileFragment does not evaluate it so
      // we will not drop "i32" even though it is not projected since we need it for
      // filtering
      auto expected_schema = schema({f64, i32});
  
      auto reader = this->GetRecordBatchReader(opts_->dataset_schema);
      auto source = this->GetFileSource(reader.get());
      auto fragment = this->MakeFragment(*source);
  
      int64_t row_count = 0;
  
      for (auto maybe_batch : PhysicalBatches(fragment)) {
        ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
        row_count += batch->num_rows();
        ASSERT_THAT(
            batch->schema()->fields(),
            ::testing::UnorderedPointwise(PointeesEqual(), expected_schema->fields()))
            << "EXPECTED:\n"
            << expected_schema->ToString() << "\nACTUAL:\n"
            << batch->schema()->ToString();
      }
  
      ASSERT_EQ(row_count, expected_rows());
    }
    void TestScanProjectedNested(bool fine_grained_selection = false) {
      auto f32 = field("f32", float32());
      auto f64 = field("f64", float64());
      auto i32 = field("i32", int32());
      auto i64 = field("i64", int64());
      auto struct1 = field("struct1", struct_({f32, i32}));
      auto struct2 = field("struct2", struct_({f64, i64, struct1}));
      this->SetSchema({struct1, struct2, f32, f64, i32, i64});
      this->ProjectNested({".struct1.f32", ".struct2.struct1", ".struct2.struct1.f32"});
      this->SetFilter(greater_equal(field_ref(FieldRef("struct2", "i64")), literal(0)));
  
      std::shared_ptr<Schema> physical_schema;
      if (fine_grained_selection) {
        // Some formats, like Parquet, let you pluck only a part of a complex type
        physical_schema = schema(
            {field("struct1", struct_({f32})), field("struct2", struct_({i64, struct1}))});
      } else {
        // Otherwise, the entire top-level field is returned
        physical_schema = schema({struct1, struct2});
      }
      std::shared_ptr<Schema> projected_schema = schema({
          field(".struct1.f32", float32()),
          field(".struct2.struct1", struct1->type()),
          field(".struct2.struct1.f32", float32()),
      });
  
      {
        auto reader = this->GetRecordBatchReader(opts_->dataset_schema);
        auto source = this->GetFileSource(reader.get());
        auto fragment = this->MakeFragment(*source);
  
        int64_t row_count = 0;
        for (auto maybe_batch : PhysicalBatches(fragment)) {
          ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
          row_count += batch->num_rows();
          ASSERT_THAT(
              batch->schema()->fields(),
              ::testing::UnorderedPointwise(PointeesEqual(), physical_schema->fields()))
              << "EXPECTED:\n"
              << physical_schema->ToString() << "\nACTUAL:\n"
              << batch->schema()->ToString();
        }
        ASSERT_EQ(row_count, expected_rows());
      }
      {
        auto reader = this->GetRecordBatchReader(opts_->dataset_schema);
        auto source = this->GetFileSource(reader.get());
        auto fragment = this->MakeFragment(*source);
  
        int64_t row_count = 0;
        for (auto maybe_batch : Batches(fragment)) {
          ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
          row_count += batch->num_rows();
          AssertSchemaEqual(*batch->schema(), *projected_schema, /*check_metadata=*/false);
        }
        ASSERT_LE(row_count, expected_rows());
        ASSERT_GT(row_count, 0);
      }
      {
        // File includes a duplicated name in struct2
        auto struct2_physical = field("struct2", struct_({f64, i64, struct1, i64}));
        auto reader = this->GetRecordBatchReader(
            schema({struct1, struct2_physical, f32, f64, i32, i64}));
        auto source = this->GetFileSource(reader.get());
        auto fragment = this->MakeFragment(*source);
  
        auto iterator = PhysicalBatches(fragment);
        EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, ::testing::HasSubstr("i64"),
                                        iterator.Next().status());
      }
      {
        // File is missing a child in struct1
        auto struct1_physical = field("struct1", struct_({i32}));
        auto reader = this->GetRecordBatchReader(
            schema({struct1_physical, struct2, f32, f64, i32, i64}));
        auto source = this->GetFileSource(reader.get());
        auto fragment = this->MakeFragment(*source);
  
        physical_schema = schema({physical_schema->field(1)});
  
        int64_t row_count = 0;
        for (auto maybe_batch : PhysicalBatches(fragment)) {
          ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
          row_count += batch->num_rows();
          ASSERT_THAT(
              batch->schema()->fields(),
              ::testing::UnorderedPointwise(PointeesEqual(), physical_schema->fields()))
              << "EXPECTED:\n"
              << physical_schema->ToString() << "\nACTUAL:\n"
              << batch->schema()->ToString();
        }
        ASSERT_EQ(row_count, expected_rows());
      }
    }
    void TestScanProjectedMissingCols() {
      auto f32 = field("f32", float32());
      auto f64 = field("f64", float64());
      auto i32 = field("i32", int32());
      auto i64 = field("i64", int64());
      this->SetSchema({f64, i64, f32, i32});
      this->Project({"f64"});
      this->SetFilter(equal(field_ref("i32"), literal(0)));
  
      auto reader_without_i32 = this->GetRecordBatchReader(schema({f64, i64, f32}));
      auto reader_without_f64 = this->GetRecordBatchReader(schema({i64, f32, i32}));
      auto reader = this->GetRecordBatchReader(schema({f64, i64, f32, i32}));
  
      auto readers = {reader.get(), reader_without_i32.get(), reader_without_f64.get()};
      for (auto reader : readers) {
        SCOPED_TRACE(reader->schema()->ToString());
        auto source = this->GetFileSource(reader);
        auto fragment = this->MakeFragment(*source);
  
        // NB: projection is applied by the scanner; FileFragment does not evaluate it so
        // we will not drop "i32" even though it is not projected since we need it for
        // filtering
        //
        // in the case where a file doesn't contain a referenced field, we won't
        // materialize it as nulls later
        std::shared_ptr<Schema> expected_schema;
        if (reader == reader_without_i32.get()) {
          expected_schema = schema({f64});
        } else if (reader == reader_without_f64.get()) {
          expected_schema = schema({i32});
        } else {
          expected_schema = schema({f64, i32});
        }
  
        int64_t row_count = 0;
        for (auto maybe_batch : PhysicalBatches(fragment)) {
          ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
          row_count += batch->num_rows();
          ASSERT_THAT(
              batch->schema()->fields(),
              ::testing::UnorderedPointwise(PointeesEqual(), expected_schema->fields()))
              << "EXPECTED:\n"
              << expected_schema->ToString() << "\nACTUAL:\n"
              << batch->schema()->ToString();
        }
        ASSERT_EQ(row_count, expected_rows());
      }
    }
    void TestScanWithVirtualColumn() {
      auto reader = this->GetRecordBatchReader(schema({field("f64", float64())}));
      auto source = this->GetFileSource(reader.get());
      // NB: dataset_schema includes a column not present in the file
      this->SetSchema({reader->schema()->field(0), field("virtual", int32())});
      auto fragment = this->MakeFragment(*source);
  
      ASSERT_OK_AND_ASSIGN(auto physical_schema, fragment->ReadPhysicalSchema());
      AssertSchemaEqual(Schema({field("f64", float64())}), *physical_schema);
      {
        int64_t row_count = 0;
        for (auto maybe_batch : Batches(fragment)) {
          ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
          AssertSchemaEqual(*batch->schema(), *opts_->projected_schema);
          row_count += batch->num_rows();
        }
        ASSERT_EQ(row_count, expected_rows());
      }
      {
        int64_t row_count = 0;
        for (auto maybe_batch : PhysicalBatches(fragment)) {
          ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
          AssertSchemaEqual(*batch->schema(), *physical_schema);
          row_count += batch->num_rows();
        }
        ASSERT_EQ(row_count, expected_rows());
      }
    }
    void TestScanWithDuplicateColumn() {
      // A duplicate column is ignored if not requested.
      auto i32 = field("i32", int32());
      auto i64 = field("i64", int64());
      this->opts_->dataset_schema = schema({i32, i32, i64});
      this->Project({"i64"});
      auto expected_schema = schema({i64});
      auto reader = this->GetRecordBatchReader(opts_->dataset_schema);
      auto source = this->GetFileSource(reader.get());
      auto fragment = this->MakeFragment(*source);
  
      int64_t row_count = 0;
  
      for (auto maybe_batch : PhysicalBatches(fragment)) {
        ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
        row_count += batch->num_rows();
        AssertSchemaEqual(*batch->schema(), *expected_schema,
                          /*check_metadata=*/false);
      }
  
      ASSERT_EQ(row_count, expected_rows());
    }
    void TestScanWithDuplicateColumnError() {
      // A duplicate column leads to an error if requested.
      auto i32 = field("i32", int32());
      auto i64 = field("i64", int64());
      this->opts_->dataset_schema = schema({i32, i32, i64});
      ASSERT_RAISES(Invalid,
                    ProjectionDescr::FromNames({"i32"}, *this->opts_->dataset_schema));
    }
    void TestScanWithPushdownNulls() {
      // Regression test for ARROW-15312
      auto i64 = field("i64", int64());
      this->SetSchema({i64});
      this->SetFilter(is_null(field_ref("i64")));
  
      auto rb = RecordBatchFromJSON(schema({i64}), R"([
        [null],
        [32]
      ])");
      ASSERT_OK_AND_ASSIGN(auto reader, RecordBatchReader::Make({rb}));
      auto source = this->GetFileSource(reader.get());
  
      auto fragment = this->MakeFragment(*source);
      int64_t row_count = 0;
      for (auto maybe_batch : Batches(fragment)) {
        ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
        row_count += batch->num_rows();
      }
      ASSERT_EQ(row_count, 1);
    }
  
   protected:
    using FileFormatFixtureMixin<FormatHelper>::opts_;
  };
  
  /// \brief A dummy FileFormat implementation
  class DummyFileFormat : public FileFormat {
   public:
    explicit DummyFileFormat(std::shared_ptr<Schema> schema = NULLPTR)
        : schema_(std::move(schema)) {}
  
    std::string type_name() const override { return "dummy"; }
  
    bool Equals(const FileFormat& other) const override {
      return type_name() == other.type_name() &&
             schema_->Equals(checked_cast<const DummyFileFormat&>(other).schema_);
    }
  
    Result<bool> IsSupported(const FileSource& source) const override { return true; }
  
    Result<std::shared_ptr<Schema>> Inspect(const FileSource& source) const override {
      return schema_;
    }
  
    /// \brief Open a file for scanning (always returns an empty generator)
    Result<RecordBatchGenerator> ScanBatchesAsync(
        const std::shared_ptr<ScanOptions>& options,
        const std::shared_ptr<FileFragment>& fragment) const override {
      return MakeEmptyGenerator<std::shared_ptr<RecordBatch>>();
    }
  
    Result<std::shared_ptr<FileWriter>> MakeWriter(
        std::shared_ptr<io::OutputStream> destination, std::shared_ptr<Schema> schema,
        std::shared_ptr<FileWriteOptions> options,
        fs::FileLocator destination_locator) const override {
      return Status::NotImplemented("writing fragment of DummyFileFormat");
    }
  
    std::shared_ptr<FileWriteOptions> DefaultWriteOptions() override { return nullptr; }
  
   protected:
    std::shared_ptr<Schema> schema_;
  };
  
  class JSONRecordBatchFileFormat : public FileFormat {
   public:
    using SchemaResolver = std::function<std::shared_ptr<Schema>(const FileSource&)>;
  
    explicit JSONRecordBatchFileFormat(std::shared_ptr<Schema> schema)
        : resolver_([schema](const FileSource&) { return schema; }) {}
  
    explicit JSONRecordBatchFileFormat(SchemaResolver resolver)
        : resolver_(std::move(resolver)) {}
  
    bool Equals(const FileFormat& other) const override { return this == &other; }
  
    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; }
  
    Result<std::shared_ptr<Schema>> Inspect(const FileSource& source) const override {
      return resolver_(source);
    }
  
    Result<RecordBatchGenerator> ScanBatchesAsync(
        const std::shared_ptr<ScanOptions>& options,
        const std::shared_ptr<FileFragment>& fragment) const override {
      ARROW_ASSIGN_OR_RAISE(auto file, fragment->source().Open());
      ARROW_ASSIGN_OR_RAISE(int64_t size, file->GetSize());
      ARROW_ASSIGN_OR_RAISE(auto buffer, file->Read(size));
      ARROW_ASSIGN_OR_RAISE(auto schema, Inspect(fragment->source()));
  
      RecordBatchVector batches{RecordBatchFromJSON(schema, util::string_view{*buffer})};
      return MakeVectorGenerator(std::move(batches));
    }
  
    Result<std::shared_ptr<FileWriter>> MakeWriter(
        std::shared_ptr<io::OutputStream> destination, std::shared_ptr<Schema> schema,
        std::shared_ptr<FileWriteOptions> options,
        fs::FileLocator destination_locator) const override {
      return Status::NotImplemented("writing fragment of JSONRecordBatchFileFormat");
    }
  
    std::shared_ptr<FileWriteOptions> DefaultWriteOptions() override { return nullptr; }
  
   protected:
    SchemaResolver resolver_;
  };
  
  struct MakeFileSystemDatasetMixin {
    std::vector<fs::FileInfo> ParsePathList(const std::string& pathlist) {
      std::vector<fs::FileInfo> infos;
  
      std::stringstream ss(pathlist);
      std::string line;
      while (std::getline(ss, line)) {
        auto start = line.find_first_not_of(" \n\r\t");
        if (start == std::string::npos) {
          continue;
        }
        line.erase(0, start);
  
        if (line.front() == '#') {
          continue;
        }
  
        if (line.back() == '/') {
          infos.push_back(fs::Dir(line));
          continue;
        }
  
        infos.push_back(fs::File(line));
      }
  
      return infos;
    }
  
    void MakeFileSystem(const std::vector<fs::FileInfo>& infos) {
      ASSERT_OK_AND_ASSIGN(fs_, fs::internal::MockFileSystem::Make(fs::kNoTime, infos));
    }
  
    void MakeFileSystem(const std::vector<std::string>& paths) {
      std::vector<fs::FileInfo> infos{paths.size()};
      std::transform(paths.cbegin(), paths.cend(), infos.begin(),
                     [](const std::string& p) { return fs::File(p); });
  
      ASSERT_OK_AND_ASSIGN(fs_, fs::internal::MockFileSystem::Make(fs::kNoTime, infos));
    }
  
    void MakeDataset(const std::vector<fs::FileInfo>& infos,
                     compute::Expression root_partition = literal(true),
                     std::vector<compute::Expression> partitions = {},
                     std::shared_ptr<Schema> s = schema({})) {
      auto n_fragments = infos.size();
      if (partitions.empty()) {
        partitions.resize(n_fragments, literal(true));
      }
  
      MakeFileSystem(infos);
      auto format = std::make_shared<DummyFileFormat>(s);
  
      std::vector<std::shared_ptr<FileFragment>> fragments;
      for (size_t i = 0; i < n_fragments; i++) {
        const auto& info = infos[i];
        if (!info.IsFile()) {
          continue;
        }
  
        ASSERT_OK_AND_ASSIGN(auto fragment,
                             format->MakeFragment({info, fs_}, partitions[i]));
        fragments.push_back(std::move(fragment));
      }
  
      ASSERT_OK_AND_ASSIGN(dataset_, FileSystemDataset::Make(s, root_partition, format, fs_,
                                                             std::move(fragments)));
    }
  
    std::shared_ptr<fs::FileSystem> fs_;
    std::shared_ptr<Dataset> dataset_;
    std::shared_ptr<ScanOptions> options_;
  };
  
  static const std::string& PathOf(const std::shared_ptr<Fragment>& fragment) {
    EXPECT_NE(fragment, nullptr);
    EXPECT_THAT(fragment->type_name(), "dummy");
    return checked_cast<const FileFragment&>(*fragment).source().path();
  }
  
  class TestFileSystemDataset : public ::testing::Test,
                                public MakeFileSystemDatasetMixin {};
  
  static std::vector<std::string> PathsOf(const FragmentVector& fragments) {
    std::vector<std::string> paths(fragments.size());
    std::transform(fragments.begin(), fragments.end(), paths.begin(), PathOf);
    return paths;
  }
  
  void AssertFilesAre(const std::shared_ptr<Dataset>& dataset,
                      std::vector<std::string> expected) {
    auto fs_dataset = checked_cast<FileSystemDataset*>(dataset.get());
    EXPECT_THAT(fs_dataset->files(), testing::UnorderedElementsAreArray(expected));
  }
  
  void AssertFragmentsAreFromPath(FragmentIterator it, std::vector<std::string> expected) {
    // Ordering is not guaranteed.
    EXPECT_THAT(PathsOf(IteratorToVector(std::move(it))),
                testing::UnorderedElementsAreArray(expected));
  }
  
  static std::vector<compute::Expression> PartitionExpressionsOf(
      const FragmentVector& fragments) {
    std::vector<compute::Expression> partition_expressions;
    std::transform(fragments.begin(), fragments.end(),
                   std::back_inserter(partition_expressions),
                   [](const std::shared_ptr<Fragment>& fragment) {
                     return fragment->partition_expression();
                   });
    return partition_expressions;
  }
  
  void AssertFragmentsHavePartitionExpressions(std::shared_ptr<Dataset> dataset,
                                               std::vector<compute::Expression> expected) {
    ASSERT_OK_AND_ASSIGN(auto fragment_it, dataset->GetFragments());
    // Ordering is not guaranteed.
    EXPECT_THAT(PartitionExpressionsOf(IteratorToVector(std::move(fragment_it))),
                testing::UnorderedElementsAreArray(expected));
  }
  
  struct ArithmeticDatasetFixture {
    static std::shared_ptr<Schema> schema() {
      return ::arrow::schema({
          field("i64", int64()),
          field("struct", struct_({
                              field("i32", int32()),
                              field("str", utf8()),
                          })),
          field("u8", uint8()),
          field("list", list(int32())),
          field("bool", boolean()),
      });
    }
  
    /// \brief Creates a single JSON record templated with n as follow.
    ///
    /// {"i64": n, "struct": {"i32": n, "str": "n"}, "u8": n "list": [n,n], "bool": n %
    /// 2},
    static std::string JSONRecordFor(int64_t n) {
      std::stringstream ss;
      auto n_i32 = static_cast<int32_t>(n);
  
      ss << "{";
      ss << "\"i64\": " << n << ", ";
      ss << "\"struct\": {";
      {
        ss << "\"i32\": " << n_i32 << ", ";
        ss << R"("str": ")" << std::to_string(n) << "\"";
      }
      ss << "}, ";
      ss << "\"u8\": " << static_cast<int32_t>(n) << ", ";
      ss << "\"list\": [" << n_i32 << ", " << n_i32 << "], ";
      ss << "\"bool\": " << (static_cast<bool>(n % 2) ? "true" : "false");
      ss << "}";
  
      return ss.str();
    }
  
    /// \brief Creates a JSON RecordBatch
    static std::string JSONRecordBatch(int64_t n) {
      DCHECK_GT(n, 0);
  
      auto record = JSONRecordFor(n);
  
      std::stringstream ss;
      ss << "[\n";
      for (int64_t i = 1; i <= n; i++) {
        if (i != 1) {
          ss << "\n,";
        }
        ss << record;
      }
      ss << "]\n";
      return ss.str();
    }
  
    static std::shared_ptr<RecordBatch> GetRecordBatch(int64_t n) {
      return RecordBatchFromJSON(ArithmeticDatasetFixture::schema(), JSONRecordBatch(n));
    }
  
    static std::unique_ptr<RecordBatchReader> GetRecordBatchReader(int64_t n) {
      DCHECK_GT(n, 0);
  
      // Functor which generates `n` RecordBatch
      struct {
        Status operator()(std::shared_ptr<RecordBatch>* out) {
          *out = i++ < count ? GetRecordBatch(i) : nullptr;
          return Status::OK();
        }
        int64_t i;
        int64_t count;
      } generator{0, n};
  
      return MakeGeneratedRecordBatch(schema(), std::move(generator));
    }
  };
  
  class WriteFileSystemDatasetMixin : public MakeFileSystemDatasetMixin {
   public:
    using PathAndContent = std::unordered_map<std::string, std::string>;
  
    void MakeSourceDataset() {
      PathAndContent source_files;
  
      source_files["/dataset/year=2018/month=01/dat0.json"] = R"([
          {"region": "NY", "model": "3", "sales": 742.0, "country": "US"},
          {"region": "NY", "model": "S", "sales": 304.125, "country": "US"},
          {"region": "NY", "model": "Y", "sales": 27.5, "country": "US"}
        ])";
      source_files["/dataset/year=2018/month=01/dat1.json"] = R"([
          {"region": "QC", "model": "3", "sales": 512, "country": "CA"},
          {"region": "QC", "model": "S", "sales": 978, "country": "CA"},
          {"region": "NY", "model": "X", "sales": 136.25, "country": "US"},
          {"region": "QC", "model": "X", "sales": 1.0, "country": "CA"},
          {"region": "QC", "model": "Y", "sales": 69, "country": "CA"}
        ])";
      source_files["/dataset/year=2019/month=01/dat0.json"] = R"([
          {"region": "CA", "model": "3", "sales": 273.5, "country": "US"},
          {"region": "CA", "model": "S", "sales": 13, "country": "US"},
          {"region": "CA", "model": "X", "sales": 54, "country": "US"},
          {"region": "QC", "model": "S", "sales": 10, "country": "CA"},
          {"region": "CA", "model": "Y", "sales": 21, "country": "US"}
        ])";
      source_files["/dataset/year=2019/month=01/dat1.json"] = R"([
          {"region": "QC", "model": "3", "sales": 152.25, "country": "CA"},
          {"region": "QC", "model": "X", "sales": 42, "country": "CA"},
          {"region": "QC", "model": "Y", "sales": 37, "country": "CA"}
        ])";
      source_files["/dataset/.pesky"] = "garbage content";
  
      auto mock_fs = std::make_shared<fs::internal::MockFileSystem>(fs::kNoTime);
      for (const auto& f : source_files) {
        ARROW_EXPECT_OK(mock_fs->CreateFile(f.first, f.second, /* recursive */ true));
      }
      fs_ = mock_fs;
  
      /// schema for the whole dataset (both source and destination)
      source_schema_ = schema({
          field("region", utf8()),
          field("model", utf8()),
          field("sales", float64()),
          field("year", int32()),
          field("month", int32()),
          field("country", utf8()),
      });
  
      /// Dummy file format for source dataset. Note that it isn't partitioned on country
      auto source_format = std::make_shared<JSONRecordBatchFileFormat>(
          SchemaFromColumnNames(source_schema_, {"region", "model", "sales", "country"}));
  
      fs::FileSelector s;
      s.base_dir = "/dataset";
      s.recursive = true;
  
      FileSystemFactoryOptions options;
      options.selector_ignore_prefixes = {"."};
      options.partitioning = std::make_shared<HivePartitioning>(
          SchemaFromColumnNames(source_schema_, {"year", "month"}));
      ASSERT_OK_AND_ASSIGN(auto factory,
                           FileSystemDatasetFactory::Make(fs_, s, source_format, options));
      ASSERT_OK_AND_ASSIGN(dataset_, factory->Finish());
  
      scan_options_ = std::make_shared<ScanOptions>();
      scan_options_->dataset_schema = dataset_->schema();
      ASSERT_OK_AND_ASSIGN(
          auto projection,
          ProjectionDescr::FromNames(source_schema_->field_names(), *dataset_->schema()));
      SetProjection(scan_options_.get(), std::move(projection));
    }
  
    void SetWriteOptions(std::shared_ptr<FileWriteOptions> file_write_options) {
      write_options_.file_write_options = file_write_options;
      write_options_.filesystem = fs_;
      write_options_.base_dir = "/new_root/";
      write_options_.basename_template = "dat_{i}";
      write_options_.writer_pre_finish = [this](FileWriter* writer) {
        visited_paths_.push_back(writer->destination().path);
        return Status::OK();
      };
    }
  
    void DoWrite(std::shared_ptr<Partitioning> desired_partitioning) {
      write_options_.partitioning = desired_partitioning;
      auto scanner_builder = ScannerBuilder(dataset_, scan_options_);
      ASSERT_OK_AND_ASSIGN(auto scanner, scanner_builder.Finish());
      ASSERT_OK(FileSystemDataset::Write(write_options_, scanner));
  
      // re-discover the written dataset
      fs::FileSelector s;
      s.recursive = true;
      s.base_dir = "/new_root";
  
      FileSystemFactoryOptions factory_options;
      factory_options.partitioning = desired_partitioning;
      ASSERT_OK_AND_ASSIGN(
          auto factory, FileSystemDatasetFactory::Make(fs_, s, format_, factory_options));
      ASSERT_OK_AND_ASSIGN(written_, factory->Finish());
    }
  
    void TestWriteWithIdenticalPartitioningSchema() {
      DoWrite(std::make_shared<DirectoryPartitioning>(
          SchemaFromColumnNames(source_schema_, {"year", "month"})));
  
      expected_files_["/new_root/2018/1/dat_0"] = R"([
          {"region": "QC", "model": "X", "sales": 1.0, "country": "CA"},
          {"region": "NY", "model": "Y", "sales": 27.5, "country": "US"},
          {"region": "QC", "model": "Y", "sales": 69, "country": "CA"},
          {"region": "NY", "model": "X", "sales": 136.25, "country": "US"},
          {"region": "NY", "model": "S", "sales": 304.125, "country": "US"},
          {"region": "QC", "model": "3", "sales": 512, "country": "CA"},
          {"region": "NY", "model": "3", "sales": 742.0, "country": "US"},
          {"region": "QC", "model": "S", "sales": 978, "country": "CA"}
        ])";
      expected_files_["/new_root/2019/1/dat_0"] = R"([
          {"region": "QC", "model": "S", "sales": 10, "country": "CA"},
          {"region": "CA", "model": "S", "sales": 13, "country": "US"},
          {"region": "CA", "model": "Y", "sales": 21, "country": "US"},
          {"region": "QC", "model": "Y", "sales": 37, "country": "CA"},
          {"region": "QC", "model": "X", "sales": 42, "country": "CA"},
          {"region": "CA", "model": "X", "sales": 54, "country": "US"},
          {"region": "QC", "model": "3", "sales": 152.25, "country": "CA"},
          {"region": "CA", "model": "3", "sales": 273.5, "country": "US"}
        ])";
      expected_physical_schema_ =
          SchemaFromColumnNames(source_schema_, {"region", "model", "sales", "country"});
  
      AssertWrittenAsExpected();
    }
  
    void TestWriteWithUnrelatedPartitioningSchema() {
      DoWrite(std::make_shared<DirectoryPartitioning>(
          SchemaFromColumnNames(source_schema_, {"country", "region"})));
  
      // XXX first thing a user will be annoyed by: we don't support left
      // padding the month field with 0.
      expected_files_["/new_root/US/NY/dat_0"] = R"([
          {"year": 2018, "month": 1, "model": "Y", "sales": 27.5},
          {"year": 2018, "month": 1, "model": "X", "sales": 136.25},
          {"year": 2018, "month": 1, "model": "S", "sales": 304.125},
          {"year": 2018, "month": 1, "model": "3", "sales": 742.0}
      ])";
      expected_files_["/new_root/CA/QC/dat_0"] = R"([
          {"year": 2018, "month": 1, "model": "X", "sales": 1.0},
          {"year": 2019, "month": 1, "model": "S", "sales": 10},
          {"year": 2019, "month": 1, "model": "Y", "sales": 37},
          {"year": 2019, "month": 1, "model": "X", "sales": 42},
          {"year": 2018, "month": 1, "model": "Y", "sales": 69},
          {"year": 2019, "month": 1, "model": "3", "sales": 152.25},
          {"year": 2018, "month": 1, "model": "3", "sales": 512},
          {"year": 2018, "month": 1, "model": "S", "sales": 978}
      ])";
      expected_files_["/new_root/US/CA/dat_0"] = R"([
          {"year": 2019, "month": 1, "model": "S", "sales": 13},
          {"year": 2019, "month": 1, "model": "Y", "sales": 21},
          {"year": 2019, "month": 1, "model": "X", "sales": 54},
          {"year": 2019, "month": 1, "model": "3", "sales": 273.5}
      ])";
      expected_physical_schema_ =
          SchemaFromColumnNames(source_schema_, {"model", "sales", "year", "month"});
  
      AssertWrittenAsExpected();
    }
  
    void TestWriteWithSupersetPartitioningSchema() {
      DoWrite(std::make_shared<DirectoryPartitioning>(
          SchemaFromColumnNames(source_schema_, {"year", "month", "country", "region"})));
  
      // XXX first thing a user will be annoyed by: we don't support left
      // padding the month field with 0.
      expected_files_["/new_root/2018/1/US/NY/dat_0"] = R"([
          {"model": "Y", "sales": 27.5},
          {"model": "X", "sales": 136.25},
          {"model": "S", "sales": 304.125},
          {"model": "3", "sales": 742.0}
      ])";
      expected_files_["/new_root/2018/1/CA/QC/dat_0"] = R"([
          {"model": "X", "sales": 1.0},
          {"model": "Y", "sales": 69},
          {"model": "3", "sales": 512},
          {"model": "S", "sales": 978}
      ])";
      expected_files_["/new_root/2019/1/US/CA/dat_0"] = R"([
          {"model": "S", "sales": 13},
          {"model": "Y", "sales": 21},
          {"model": "X", "sales": 54},
          {"model": "3", "sales": 273.5}
      ])";
      expected_files_["/new_root/2019/1/CA/QC/dat_0"] = R"([
          {"model": "S", "sales": 10},
          {"model": "Y", "sales": 37},
          {"model": "X", "sales": 42},
          {"model": "3", "sales": 152.25}
      ])";
      expected_physical_schema_ = SchemaFromColumnNames(source_schema_, {"model", "sales"});
  
      AssertWrittenAsExpected();
    }
  
    void TestWriteWithEmptyPartitioningSchema() {
      DoWrite(std::make_shared<DirectoryPartitioning>(
          SchemaFromColumnNames(source_schema_, {})));
  
      expected_files_["/new_root/dat_0"] = R"([
          {"country": "CA", "region": "QC", "year": 2018, "month": 1, "model": "X", "sales": 1.0},
          {"country": "CA", "region": "QC", "year": 2019, "month": 1, "model": "S", "sales": 10},
          {"country": "US", "region": "CA", "year": 2019, "month": 1, "model": "S", "sales": 13},
          {"country": "US", "region": "CA", "year": 2019, "month": 1, "model": "Y", "sales": 21},
          {"country": "US", "region": "NY", "year": 2018, "month": 1, "model": "Y", "sales": 27.5},
          {"country": "CA", "region": "QC", "year": 2019, "month": 1, "model": "Y", "sales": 37},
          {"country": "CA", "region": "QC", "year": 2019, "month": 1, "model": "X", "sales": 42},
          {"country": "US", "region": "CA", "year": 2019, "month": 1, "model": "X", "sales": 54},
          {"country": "CA", "region": "QC", "year": 2018, "month": 1, "model": "Y", "sales": 69},
          {"country": "US", "region": "NY", "year": 2018, "month": 1, "model": "X", "sales": 136.25},
          {"country": "CA", "region": "QC", "year": 2019, "month": 1, "model": "3", "sales": 152.25},
          {"country": "US", "region": "CA", "year": 2019, "month": 1, "model": "3", "sales": 273.5},
          {"country": "US", "region": "NY", "year": 2018, "month": 1, "model": "S", "sales": 304.125},
          {"country": "CA", "region": "QC", "year": 2018, "month": 1, "model": "3", "sales": 512},
          {"country": "US", "region": "NY", "year": 2018, "month": 1, "model": "3", "sales": 742.0},
          {"country": "CA", "region": "QC", "year": 2018, "month": 1, "model": "S", "sales": 978}
      ])";
      expected_physical_schema_ = source_schema_;
  
      AssertWrittenAsExpected();
    }
  
    void AssertWrittenAsExpected() {
      std::unordered_set<std::string> expected_paths, actual_paths;
      for (const auto& file_contents : expected_files_) {
        expected_paths.insert(file_contents.first);
      }
  
      // expect the written filesystem to contain precisely the paths we expected
      for (auto path : checked_pointer_cast<FileSystemDataset>(written_)->files()) {
        actual_paths.insert(std::move(path));
      }
      EXPECT_THAT(actual_paths, testing::UnorderedElementsAreArray(expected_paths));
  
      // Additionally, the writer producing each written file was visited and its path
      // collected. That should match the expected paths as well
      EXPECT_THAT(visited_paths_, testing::UnorderedElementsAreArray(expected_paths));
  
      ASSERT_OK_AND_ASSIGN(auto written_fragments_it, written_->GetFragments());
      for (auto maybe_fragment : written_fragments_it) {
        ASSERT_OK_AND_ASSIGN(auto fragment, maybe_fragment);
  
        ASSERT_OK_AND_ASSIGN(auto actual_physical_schema, fragment->ReadPhysicalSchema());
        AssertSchemaEqual(*expected_physical_schema_, *actual_physical_schema,
                          check_metadata_);
  
        const auto& path = checked_pointer_cast<FileFragment>(fragment)->source().path();
  
        auto file_contents = expected_files_.find(path);
        if (file_contents == expected_files_.end()) {
          // file wasn't expected to be written at all; nothing to compare with
          continue;
        }
  
        ASSERT_OK_AND_ASSIGN(auto scanner, ScannerBuilder(actual_physical_schema, fragment,
                                                          std::make_shared<ScanOptions>())
                                               .Finish());
        ASSERT_OK_AND_ASSIGN(auto actual_table, scanner->ToTable());
        ASSERT_OK_AND_ASSIGN(actual_table, actual_table->CombineChunks());
        std::shared_ptr<Array> actual_struct;
  
        for (auto maybe_batch :
             MakeIteratorFromReader(std::make_shared<TableBatchReader>(*actual_table))) {
          ASSERT_OK_AND_ASSIGN(auto batch, maybe_batch);
          ASSERT_OK_AND_ASSIGN(
              auto sort_indices,
              compute::SortIndices(batch->GetColumnByName("sales"),
                                   compute::SortOptions({compute::SortKey{"sales"}})));
          ASSERT_OK_AND_ASSIGN(Datum sorted_batch, compute::Take(batch, sort_indices));
          ASSERT_OK_AND_ASSIGN(auto struct_array,
                               sorted_batch.record_batch()->ToStructArray());
          actual_struct = std::dynamic_pointer_cast<Array>(struct_array);
        }
  
        auto expected_struct = ArrayFromJSON(struct_(expected_physical_schema_->fields()),
                                             {file_contents->second});
  
        AssertArraysEqual(*expected_struct, *actual_struct, /*verbose=*/true);
      }
    }
  
    bool check_metadata_ = true;
    std::shared_ptr<Schema> source_schema_;
    std::shared_ptr<FileFormat> format_;
    PathAndContent expected_files_;
    std::shared_ptr<Schema> expected_physical_schema_;
    std::shared_ptr<Dataset> written_;
    std::vector<std::string> visited_paths_;
    FileSystemDatasetWriteOptions write_options_;
    std::shared_ptr<ScanOptions> scan_options_;
  };
  
  }  // namespace dataset
  }  // namespace arrow
  