Apply patches from @bkietz

cpp/src/arrow/dataset/dataset.h

@@ -94,6 +94,8 @@ class ARROW_DS_EXPORT DataSource {
   /// controls filtering and schema inference.
   DataFragmentIterator GetFragments(std::shared_ptr<ScanOptions> options);
 
+  const std::shared_ptr<Schema>& schema() const { return schema_; }
+
   /// \brief An expression which evaluates to true for all data viewed by this DataSource.
   /// May be null, which indicates no information is available.
   const std::shared_ptr<Expression>& partition_expression() const {
@@ -106,9 +108,10 @@ class ARROW_DS_EXPORT DataSource {
   virtual ~DataSource() = default;
 
  protected:
-  DataSource() = default;
-  explicit DataSource(std::shared_ptr<Expression> c)
-      : partition_expression_(std::move(c)) {}
+  explicit DataSource(std::shared_ptr<Schema> schema) : schema_(std::move(schema)) {}
+
+  DataSource(std::shared_ptr<Schema> schema, std::shared_ptr<Expression> e)
+      : partition_expression_(std::move(e)) {}
 
   virtual DataFragmentIterator GetFragmentsImpl(std::shared_ptr<ScanOptions> options) = 0;
 
@@ -118,14 +121,15 @@ class ARROW_DS_EXPORT DataSource {
       const std::shared_ptr<ScanOptions>& scan_options,
       std::shared_ptr<ScanOptions>* simplified_scan_options) const;
 
+  std::shared_ptr<Schema> schema_;
   std::shared_ptr<Expression> partition_expression_;
 };
 
 /// \brief A DataSource consisting of a flat sequence of DataFragments
 class ARROW_DS_EXPORT SimpleDataSource : public DataSource {
  public:
-  explicit SimpleDataSource(DataFragmentVector fragments)
-      : fragments_(std::move(fragments)) {}
+  explicit SimpleDataSource(std::shared_ptr<Schema> schema, DataFragmentVector fragments)
+      : DataSource(std::move(schema)), fragments_(std::move(fragments)) {}
 
   DataFragmentIterator GetFragmentsImpl(std::shared_ptr<ScanOptions> options) override;
 
@@ -138,7 +142,8 @@ class ARROW_DS_EXPORT SimpleDataSource : public DataSource {
 /// \brief A recursive DataSource with child DataSources.
 class ARROW_DS_EXPORT TreeDataSource : public DataSource {
  public:
-  explicit TreeDataSource(DataSourceVector children) : children_(std::move(children)) {}
+  explicit TreeDataSource(std::shared_ptr<Schema> schema, DataSourceVector children)
+      : DataSource(std::move(schema)), children_(std::move(children)) {}
 
   DataFragmentIterator GetFragmentsImpl(std::shared_ptr<ScanOptions> options) override;
 

cpp/src/arrow/dataset/dataset_test.cc

@@ -62,7 +62,7 @@ TEST_F(TestSimpleDataSource, GetFragments) {
   auto fragment = std::make_shared<SimpleDataFragment>(batches, options_);
   // It is safe to copy fragment multiple time since Scan() does not consume
   // the internal array.
-  auto source = SimpleDataSource({kNumberFragments, fragment});
+  auto source = SimpleDataSource(schema_, {kNumberFragments, fragment});
 
   AssertDataSourceEquals(reader.get(), &source);
 }
@@ -87,17 +87,17 @@ TEST_F(TestTreeDataSource, GetFragments) {
   // Creates a complete binary tree of depth kCompleteBinaryTreeDepth where the
   // leaves are SimpleDataSource containing kChildPerNode fragments.
 
-  auto l1_leaf_source =
-      std::make_shared<SimpleDataSource>(DataFragmentVector{kChildPerNode, fragment});
+  auto l1_leaf_source = std::make_shared<SimpleDataSource>(
+      schema_, DataFragmentVector{kChildPerNode, fragment});
 
-  auto l2_leaf_tree_source =
-      std::make_shared<TreeDataSource>(DataSourceVector{kChildPerNode, l1_leaf_source});
+  auto l2_leaf_tree_source = std::make_shared<TreeDataSource>(
+      schema_, DataSourceVector{kChildPerNode, l1_leaf_source});
 
   auto l3_middle_tree_source = std::make_shared<TreeDataSource>(
-      DataSourceVector{kChildPerNode, l2_leaf_tree_source});
+      schema_, DataSourceVector{kChildPerNode, l2_leaf_tree_source});
 
   auto root_source = std::make_shared<TreeDataSource>(
-      DataSourceVector{kChildPerNode, l3_middle_tree_source});
+      schema_, DataSourceVector{kChildPerNode, l3_middle_tree_source});
 
   AssertDataSourceEquals(reader.get(), root_source.get());
 }
@@ -117,8 +117,8 @@ TEST_F(TestDataset, TrivialScan) {
   DataFragmentVector fragments{kNumberFragments, fragment};
 
   DataSourceVector sources = {
-      std::make_shared<SimpleDataSource>(fragments),
-      std::make_shared<SimpleDataSource>(fragments),
+      std::make_shared<SimpleDataSource>(schema_, fragments),
+      std::make_shared<SimpleDataSource>(schema_, fragments),
   };
 
   const int64_t total_batches = sources.size() * kNumberFragments * kNumberBatches;

cpp/src/arrow/dataset/discovery.cc

@@ -217,8 +217,8 @@ Result<std::shared_ptr<DataSource>> FileSystemDataSourceDiscovery::Finish(
 
   RETURN_NOT_OK(forest_.Visit(apply_partition_scheme));
 
-  return FileSystemDataSource::Make(fs_, forest_, std::move(partitions), root_partition_,
-                                    format_);
+  return FileSystemDataSource::Make(schema, root_partition_, format_, fs_, forest_,
+                                    std::move(partitions));
 }
 
 }  // namespace dataset

cpp/src/arrow/dataset/discovery_test.cc

@@ -58,9 +58,10 @@ class MockDataSourceDiscovery : public DataSourceDiscovery {
     return schemas_;
   }
 
-  Result<std::shared_ptr<DataSource>> Finish(const std::shared_ptr<Schema>&) override {
+  Result<std::shared_ptr<DataSource>> Finish(
+      const std::shared_ptr<Schema>& schema) override {
     return std::make_shared<SimpleDataSource>(
-        std::vector<std::shared_ptr<DataFragment>>{});
+        schema, std::vector<std::shared_ptr<DataFragment>>{});
   }
 
  protected:

cpp/src/arrow/dataset/file_base.cc

@@ -43,43 +43,45 @@ Result<ScanTaskIterator> FileDataFragment::Scan(std::shared_ptr<ScanContext> con
   return format_->ScanFile(source_, scan_options_, context);
 }
 
-FileSystemDataSource::FileSystemDataSource(std::shared_ptr<fs::FileSystem> filesystem,
+FileSystemDataSource::FileSystemDataSource(std::shared_ptr<Schema> schema,
+                                           std::shared_ptr<Expression> root_partition,
+                                           std::shared_ptr<FileFormat> format,
+                                           std::shared_ptr<fs::FileSystem> filesystem,
                                            fs::PathForest forest,
-                                           ExpressionVector file_partitions,
-                                           std::shared_ptr<Expression> source_partition,
-                                           std::shared_ptr<FileFormat> format)
-    : DataSource(std::move(source_partition)),
+                                           ExpressionVector file_partitions)
+    : DataSource(std::move(schema), std::move(root_partition)),
+      format_(std::move(format)),
       filesystem_(std::move(filesystem)),
       forest_(std::move(forest)),
-      partitions_(std::move(file_partitions)),
-      format_(std::move(format)) {
+      partitions_(std::move(file_partitions)) {
   DCHECK_EQ(static_cast<size_t>(forest_.size()), partitions_.size());
 }
 
 Result<std::shared_ptr<DataSource>> FileSystemDataSource::Make(
-    std::shared_ptr<fs::FileSystem> filesystem, fs::FileStatsVector stats,
-    std::shared_ptr<Expression> source_partition, std::shared_ptr<FileFormat> format) {
+    std::shared_ptr<Schema> schema, std::shared_ptr<Expression> root_partition,
+    std::shared_ptr<FileFormat> format, std::shared_ptr<fs::FileSystem> filesystem,
+    fs::FileStatsVector stats) {
   ExpressionVector partitions(stats.size(), scalar(true));
-  return Make(std::move(filesystem), std::move(stats), std::move(partitions),
-              std::move(source_partition), std::move(format));
+  return Make(std::move(schema), std::move(root_partition), std::move(format),
+              std::move(filesystem), std::move(stats), std::move(partitions));
 }
 
 Result<std::shared_ptr<DataSource>> FileSystemDataSource::Make(
-    std::shared_ptr<fs::FileSystem> filesystem, fs::FileStatsVector stats,
-    ExpressionVector partitions, std::shared_ptr<Expression> source_partition,
-    std::shared_ptr<FileFormat> format) {
+    std::shared_ptr<Schema> schema, std::shared_ptr<Expression> root_partition,
+    std::shared_ptr<FileFormat> format, std::shared_ptr<fs::FileSystem> filesystem,
+    fs::FileStatsVector stats, ExpressionVector partitions) {
   ARROW_ASSIGN_OR_RAISE(auto forest, fs::PathForest::Make(std::move(stats), &partitions));
-  return Make(std::move(filesystem), std::move(forest), std::move(partitions),
-              std::move(source_partition), std::move(format));
+  return Make(std::move(schema), std::move(root_partition), std::move(format),
+              std::move(filesystem), std::move(forest), std::move(partitions));
 }
 
 Result<std::shared_ptr<DataSource>> FileSystemDataSource::Make(
-    std::shared_ptr<fs::FileSystem> filesystem, fs::PathForest forest,
-    ExpressionVector partitions, std::shared_ptr<Expression> source_partition,
-    std::shared_ptr<FileFormat> format) {
+    std::shared_ptr<Schema> schema, std::shared_ptr<Expression> root_partition,
+    std::shared_ptr<FileFormat> format, std::shared_ptr<fs::FileSystem> filesystem,
+    fs::PathForest forest, ExpressionVector partitions) {
   return std::shared_ptr<DataSource>(new FileSystemDataSource(
-      std::move(filesystem), std::move(forest), std::move(partitions),
-      std::move(source_partition), std::move(format)));
+      std::move(schema), std::move(root_partition), std::move(format),
+      std::move(filesystem), std::move(forest), std::move(partitions)));
 }
 
 std::string FileSystemDataSource::ToString() const {

cpp/src/arrow/dataset/file_base.h

@@ -54,8 +54,8 @@ class ARROW_DS_EXPORT FileSource {
       : impl_(PathAndFileSystem{std::move(path), filesystem}),
         compression_(compression) {}
 
-  FileSource(std::shared_ptr<Buffer> buffer,
-             Compression::type compression = Compression::UNCOMPRESSED)
+  explicit FileSource(std::shared_ptr<Buffer> buffer,
+                      Compression::type compression = Compression::UNCOMPRESSED)
       : impl_(std::move(buffer)), compression_(compression) {}
 
   bool operator==(const FileSource& other) const {
@@ -153,16 +153,16 @@ class ARROW_DS_EXPORT FileFormat {
 /// \brief A DataFragment that is stored in a file with a known format
 class ARROW_DS_EXPORT FileDataFragment : public DataFragment {
  public:
-  FileDataFragment(const FileSource& source, std::shared_ptr<FileFormat> format,
+  FileDataFragment(FileSource source, std::shared_ptr<FileFormat> format,
                    std::shared_ptr<ScanOptions> scan_options)
       : DataFragment(std::move(scan_options)),
-        source_(source),
+        source_(std::move(source)),
         format_(std::move(format)) {}
 
   Result<ScanTaskIterator> Scan(std::shared_ptr<ScanContext> context) override;
 
   const FileSource& source() const { return source_; }
-  std::shared_ptr<FileFormat> format() const { return format_; }
+  const std::shared_ptr<FileFormat>& format() const { return format_; }
 
  protected:
   FileSource source_;
@@ -174,66 +174,73 @@ class ARROW_DS_EXPORT FileSystemDataSource : public DataSource {
  public:
   /// \brief Create a FileSystemDataSource.
   ///
+  /// \param[in] schema the top-level schema of the DataSource
+  /// \param[in] root_partition the top-level partition of the DataSource
+  /// \param[in] format file format to create fragments from.
   /// \param[in] filesystem the filesystem which files are from.
   /// \param[in] stats a list of files/directories to consume.
-  /// \param[in] source_partition the top-level partition of the DataSource
   /// attach additional partition expressions to FileStats found in `stats`.
-  /// \param[in] format file format to create fragments from.
   ///
   /// The caller is not required to provide a complete coverage of nodes and
   /// partitions.
   static Result<std::shared_ptr<DataSource>> Make(
-      std::shared_ptr<fs::FileSystem> filesystem, fs::FileStatsVector stats,
-      std::shared_ptr<Expression> source_partition, std::shared_ptr<FileFormat> format);
+      std::shared_ptr<Schema> schema, std::shared_ptr<Expression> root_partition,
+      std::shared_ptr<FileFormat> format, std::shared_ptr<fs::FileSystem> filesystem,
+      fs::FileStatsVector stats);
 
   /// \brief Create a FileSystemDataSource with file-level partitions.
   ///
+  /// \param[in] schema the top-level schema of the DataSource
+  /// \param[in] root_partition the top-level partition of the DataSource
+  /// \param[in] format file format to create fragments from.
   /// \param[in] filesystem the filesystem which files are from.
   /// \param[in] stats a list of files/directories to consume.
   /// \param[in] partitions partition information associated with `stats`.
-  /// \param[in] source_partition the top-level partition of the DataSource
   /// attach additional partition expressions to FileStats found in `stats`.
-  /// \param[in] format file format to create fragments from.
   ///
   /// The caller is not required to provide a complete coverage of nodes and
   /// partitions.
   static Result<std::shared_ptr<DataSource>> Make(
-      std::shared_ptr<fs::FileSystem> filesystem, fs::FileStatsVector stats,
-      ExpressionVector partitions, std::shared_ptr<Expression> source_partition,
-      std::shared_ptr<FileFormat> format);
+      std::shared_ptr<Schema> schema, std::shared_ptr<Expression> root_partition,
+      std::shared_ptr<FileFormat> format, std::shared_ptr<fs::FileSystem> filesystem,
+      fs::FileStatsVector stats, ExpressionVector partitions);
 
   /// \brief Create a FileSystemDataSource with file-level partitions.
   ///
+  /// \param[in] schema the top-level schema of the DataSource
+  /// \param[in] root_partition the top-level partition of the DataSource
+  /// \param[in] format file format to create fragments from.
   /// \param[in] filesystem the filesystem which files are from.
   /// \param[in] forest a PathForest of files/directories to consume.
   /// \param[in] partitions partition information associated with `forest`.
-  /// \param[in] source_partition the top-level partition of the DataSource
   /// attach additional partition expressions to FileStats found in `forest`.
-  /// \param[in] format file format to create fragments from.
   ///
   /// The caller is not required to provide a complete coverage of nodes and
   /// partitions.
   static Result<std::shared_ptr<DataSource>> Make(
-      std::shared_ptr<fs::FileSystem> filesystem, fs::PathForest forest,
-      ExpressionVector partitions, std::shared_ptr<Expression> source_partition,
-      std::shared_ptr<FileFormat> format);
+      std::shared_ptr<Schema> schema, std::shared_ptr<Expression> root_partition,
+      std::shared_ptr<FileFormat> format, std::shared_ptr<fs::FileSystem> filesystem,
+      fs::PathForest forest, ExpressionVector partitions);
 
   std::string type_name() const override { return "filesystem"; }
 
   std::string ToString() const;
 
+  const std::shared_ptr<FileFormat>& format() const { return format_; }
+
  protected:
   DataFragmentIterator GetFragmentsImpl(std::shared_ptr<ScanOptions> options) override;
 
-  FileSystemDataSource(std::shared_ptr<fs::FileSystem> filesystem, fs::PathForest forest,
-                       ExpressionVector file_partitions,
-                       std::shared_ptr<Expression> source_partition,
-                       std::shared_ptr<FileFormat> format);
+  FileSystemDataSource(std::shared_ptr<Schema> schema,
+                       std::shared_ptr<Expression> root_partition,
+                       std::shared_ptr<FileFormat> format,
+                       std::shared_ptr<fs::FileSystem> filesystem, fs::PathForest forest,
+                       ExpressionVector file_partitions);
 
+  std::shared_ptr<FileFormat> format_;
   std::shared_ptr<fs::FileSystem> filesystem_;
   fs::PathForest forest_;
   ExpressionVector partitions_;
-  std::shared_ptr<FileFormat> format_;
 };
 
 }  // namespace dataset

cpp/src/arrow/dataset/file_ipc.cc

@@ -22,6 +22,8 @@
 #include <utility>
 #include <vector>
 
+#include <iostream>
+
 #include "arrow/dataset/dataset_internal.h"
 #include "arrow/dataset/filter.h"
 #include "arrow/dataset/scanner.h"
@@ -33,13 +35,20 @@
 namespace arrow {
 namespace dataset {
 
-Status WrapErrorWithSource(Status status, const FileSource& source) {
-  if (status.ok()) {
-    return Status::OK();
+Result<std::shared_ptr<ipc::RecordBatchFileReader>> OpenReader(
+    const FileSource& source, std::shared_ptr<io::RandomAccessFile> input = nullptr) {
+  if (input == nullptr) {
+    ARROW_ASSIGN_OR_RAISE(input, source.Open());
   }
 
-  return status.WithMessage("Could not open IPC input source '", source.path(),
-                            "': ", status.message());
+  std::shared_ptr<ipc::RecordBatchFileReader> reader;
+  auto status = ipc::RecordBatchFileReader::Open(std::move(input), &reader);
+  if (!status.ok()) {
+    return status.WithMessage("Could not open IPC input source '", source.path(),
+                              "': ", status.message());
+  }
+
+  return reader;
 }
 
 /// \brief A ScanTask backed by an Ipc file.
@@ -50,11 +59,24 @@ class IpcScanTask : public ScanTask {
       : ScanTask(std::move(options), std::move(context)), source_(std::move(source)) {}
 
   Result<RecordBatchIterator> Execute() override {
-    ARROW_ASSIGN_OR_RAISE(auto input, source_.Open());
-    std::shared_ptr<RecordBatchReader> reader;
-    RETURN_NOT_OK(
-        WrapErrorWithSource(ipc::RecordBatchStreamReader::Open(input, &reader), source_));
-    return MakeFunctionIterator([reader] { return reader->Next(); });
+    struct {
+      Result<std::shared_ptr<RecordBatch>> Next() {
+        if (i_ == reader_->num_record_batches()) {
+          return nullptr;
+        }
+
+        std::shared_ptr<RecordBatch> batch;
+        RETURN_NOT_OK(reader_->ReadRecordBatch(i_++, &batch));
+        return batch;
+      }
+
+      std::shared_ptr<ipc::RecordBatchFileReader> reader_;
+      int i_ = 0;
+    } batch_it;
+
+    ARROW_ASSIGN_OR_RAISE(batch_it.reader_, OpenReader(source_));
+
+    return RecordBatchIterator(std::move(batch_it));
   }
 
  private:
@@ -95,18 +117,12 @@ class IpcScanTaskIterator {
 
 Result<bool> IpcFileFormat::IsSupported(const FileSource& source) const {
   ARROW_ASSIGN_OR_RAISE(auto input, source.Open());
-  ipc::DictionaryMemo dictionary_memo;
-  std::shared_ptr<Schema> schema;
-  return ipc::ReadSchema(input.get(), &dictionary_memo, &schema).ok();
+  return OpenReader(source, input).ok();
 }
 
 Result<std::shared_ptr<Schema>> IpcFileFormat::Inspect(const FileSource& source) const {
-  ARROW_ASSIGN_OR_RAISE(auto input, source.Open());
-  ipc::DictionaryMemo dictionary_memo;
-  std::shared_ptr<Schema> schema;
-  RETURN_NOT_OK(WrapErrorWithSource(
-      ipc::ReadSchema(input.get(), &dictionary_memo, &schema), source));
-  return schema;
+  ARROW_ASSIGN_OR_RAISE(auto reader, OpenReader(source));
+  return reader->schema();
 }
 
 Result<ScanTaskIterator> IpcFileFormat::ScanFile(

cpp/src/arrow/dataset/file_ipc_test.cc

@@ -45,8 +45,9 @@ class ArrowIpcWriterMixin : public ::testing::Test {
     EXPECT_OK_AND_ASSIGN(auto sink, io::BufferOutputStream::Create(
                                         kDefaultOutputStreamSize, default_memory_pool()));
     auto writer_schema = readers[0]->schema();
+
     EXPECT_OK_AND_ASSIGN(auto writer,
-                         ipc::RecordBatchStreamWriter::Open(sink.get(), writer_schema));
+                         ipc::RecordBatchFileWriter::Open(sink.get(), writer_schema));
 
     for (auto reader : readers) {
       std::vector<std::shared_ptr<RecordBatch>> batches;
@@ -57,6 +58,8 @@ class ArrowIpcWriterMixin : public ::testing::Test {
       }
     }
 
+    ARROW_EXPECT_OK(writer->Close());
+
     EXPECT_OK_AND_ASSIGN(auto out, sink->Finish());
     return out;
   }
@@ -68,9 +71,14 @@ class ArrowIpcWriterMixin : public ::testing::Test {
   std::shared_ptr<Buffer> Write(const Table& table) {
     EXPECT_OK_AND_ASSIGN(auto sink, io::BufferOutputStream::Create(
                                         kDefaultOutputStreamSize, default_memory_pool()));
+
     EXPECT_OK_AND_ASSIGN(auto writer,
-                         ipc::RecordBatchStreamWriter::Open(sink.get(), table.schema()));
+                         ipc::RecordBatchFileWriter::Open(sink.get(), table.schema()));
+
     ARROW_EXPECT_OK(writer->WriteTable(table));
+
+    ARROW_EXPECT_OK(writer->Close());
+
     EXPECT_OK_AND_ASSIGN(auto out, sink->Finish());
     return out;
   }
@@ -134,7 +142,7 @@ TEST_F(TestIpcFileFormat, OpenFailureWithRelevantError) {
   auto format = IpcFileFormat();
 
   std::shared_ptr<Buffer> buf = std::make_shared<Buffer>(util::string_view(""));
-  auto result = format.Inspect({buf});
+  auto result = format.Inspect(FileSource(buf));
   EXPECT_RAISES_WITH_MESSAGE_THAT(Invalid, testing::HasSubstr("<Buffer>"),
                                   result.status());
 

cpp/src/arrow/dataset/file_parquet_test.cc

@@ -188,7 +188,7 @@ TEST_F(TestParquetFileFormat, OpenFailureWithRelevantError) {
   auto format = ParquetFileFormat();
 
   std::shared_ptr<Buffer> buf = std::make_shared<Buffer>(util::string_view(""));
-  auto result = format.Inspect({buf});
+  auto result = format.Inspect(FileSource(buf));
   EXPECT_RAISES_WITH_MESSAGE_THAT(IOError, testing::HasSubstr("<Buffer>"),
                                   result.status());
 

cpp/src/arrow/dataset/scanner_test.cc

@@ -40,8 +40,8 @@ class TestScanner : public DatasetFixtureMixin {
     DataFragmentVector fragments{kNumberFragments,
                                  std::make_shared<SimpleDataFragment>(batches, options_)};
 
-    DataSourceVector sources{kNumberSources,
-                             std::make_shared<SimpleDataSource>(fragments)};
+    DataSourceVector sources{
+        kNumberSources, std::make_shared<SimpleDataSource>(batch->schema(), fragments)};
 
     return Scanner{sources, options_, ctx_};
   }

cpp/src/arrow/dataset/test_util.h

@@ -294,8 +294,9 @@ class TestFileSystemDataSource : public ::testing::Test {
 
     MakeFileSystem(stats);
     auto format = std::make_shared<DummyFileFormat>();
-    ASSERT_OK_AND_ASSIGN(source_, FileSystemDataSource::Make(fs_, stats, partitions,
-                                                             source_partition, format));
+    ASSERT_OK_AND_ASSIGN(
+        source_, FileSystemDataSource::Make(schema({}), source_partition, format, fs_,
+                                            stats, partitions));
   }
 
  protected:

python/pyarrow/_dataset.pyx

@@ -581,6 +581,13 @@ cdef class DataSource:
     cdef shared_ptr[CDataSource] unwrap(self) nogil:
         return self.wrapped
 
+    @property
+    def schema(self):
+        """
+        Schema of all RecordBatches contained in this DataSource.
+        """
+        return pyarrow_wrap_schema(self.source.schema())
+
     @property
     def partition_expression(self):
         """
@@ -601,7 +608,7 @@ cdef class TreeDataSource(DataSource):
     cdef:
         CTreeDataSource* tree_source
 
-    def __init__(self, data_sources):
+    def __init__(self, schema, data_sources):
         cdef:
             DataSource child
             CDataSourceVector children
@@ -610,7 +617,8 @@ cdef class TreeDataSource(DataSource):
         for child in data_sources:
             children.push_back(child.wrapped)
 
-        tree_source = make_shared[CTreeDataSource](children)
+        tree_source = make_shared[CTreeDataSource](
+            pyarrow_unwrap_schema(schema), children)
         self.init(<shared_ptr[CDataSource]> tree_source)
 
     cdef void init(self, const shared_ptr[CDataSource]& sp):
@@ -624,27 +632,29 @@ cdef class FileSystemDataSource(DataSource):
     cdef:
         CFileSystemDataSource* filesystem_source
 
-    def __init__(self, FileSystem filesystem not None, paths_or_selector,
-                 partitions, Expression source_partition,
-                 FileFormat file_format not None):
+    def __init__(self, Schema schema not None, Expression source_partition not None,
+                 FileFormat file_format not None, FileSystem filesystem not None,
+                 paths_or_selector, partitions):
         """Create a FileSystemDataSource
 
         Parameters
         ----------
+        schema : Schema
+            Schema for resulting DataSource
+        source_partition : Expression
+        file_format : FileFormat
         filesystem : FileSystem
-            Filesystem to discover.
+            FileSystem which will be explored to discover data files.
         paths_or_selector : Union[FileSelector, List[FileStats]]
             The file stats object can be queried by the
             filesystem.get_target_stats method.
         partitions : List[Expression]
-        source_partition : Expression
-        file_format : FileFormat
         """
         cdef:
+            shared_ptr[CExpression] c_source_partition
             FileStats stats
             Expression expression
             vector[CFileStats] c_file_stats
-            shared_ptr[CExpression] c_source_partition
             vector[shared_ptr[CExpression]] c_partitions
             CResult[shared_ptr[CDataSource]] result
 
@@ -664,11 +674,12 @@ cdef class FileSystemDataSource(DataSource):
             c_source_partition = source_partition.unwrap()
 
         result = CFileSystemDataSource.Make(
+            pyarrow_unwrap_schema(schema),
+            c_source_partition,
+            file_format.unwrap(),
             filesystem.unwrap(),
             c_file_stats,
-            c_partitions,
-            c_source_partition,
-            file_format.unwrap()
+            c_partitions
         )
         self.init(GetResultValue(result))
 

python/pyarrow/includes/libarrow_dataset.pxd

@@ -187,6 +187,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
 
     cdef cppclass CDataSource "arrow::dataset::DataSource":
         CDataFragmentIterator GetFragments(shared_ptr[CScanOptions] options)
+        const shared_ptr[CSchema]& schema()
         const shared_ptr[CExpression]& partition_expression()
         c_string type_name()
 
@@ -258,11 +259,12 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
             "arrow::dataset::FileSystemDataSource"(CDataSource):
         @staticmethod
         CResult[shared_ptr[CDataSource]] Make(
+            shared_ptr[CSchema] schema,
+            shared_ptr[CExpression] source_partition,
+            shared_ptr[CFileFormat] format,
             shared_ptr[CFileSystem] filesystem,
             CFileStatsVector stats,
-            CExpressionVector partitions,
-            shared_ptr[CExpression] source_partition,
-            shared_ptr[CFileFormat] format)
+            CExpressionVector partitions)
         c_string type()
         shared_ptr[CDataFragmentIterator] GetFragments(
             shared_ptr[CScanOptions] options)

python/pyarrow/ipc.py

@@ -69,8 +69,8 @@ sink : str, pyarrow.NativeFile, or file-like Python object
 schema : pyarrow.Schema
     The Arrow schema for data to be written to the file
 use_legacy_format : boolean, default None
-    If None, use True unless overridden by ARROW_PRE_0_15_IPC_FORMAT=1
-    environment variable"""
+    If None, False will be used unless this default is overridden by
+    setting the environment variable ARROW_PRE_0_15_IPC_FORMAT=1"""
 
 
 class RecordBatchStreamWriter(lib._RecordBatchStreamWriter):

python/pyarrow/tests/test_dataset.py

@@ -81,14 +81,15 @@ def dataset(mockfs):
 
 
 def test_filesystem_data_source(mockfs):
+    schema = pa.schema([])
+
     file_format = ds.ParquetFileFormat()
 
     paths = ['subdir/1/xxx/file0.parquet', 'subdir/2/yyy/file1.parquet']
     partitions = [ds.ScalarExpression(True), ds.ScalarExpression(True)]
 
-    source = ds.FileSystemDataSource(mockfs, paths, partitions,
-                                     source_partition=None,
-                                     file_format=file_format)
+    source = ds.FileSystemDataSource(schema,
+                                     source_partition=None, file_format=file_format, filesystem=mockfs, paths_or_selector=paths, partitions=partitions)
 
     source_partition = ds.ComparisonExpression(
         ds.CompareOperator.Equal,
@@ -107,8 +108,9 @@ def test_filesystem_data_source(mockfs):
             ds.ScalarExpression(2)
         )
     ]
-    source = ds.FileSystemDataSource(mockfs, paths, partitions,
+    source = ds.FileSystemDataSource(paths_or_selector=paths, schema=schema,
                                      source_partition=source_partition,
+                                     filesystem=mockfs, partitions=partitions,
                                      file_format=file_format)
     assert source.partition_expression.equals(source_partition)
 

r/R/arrowExports.R

@@ -368,6 +368,10 @@ dataset___DSDiscovery__Finish2 <- function(discovery, schema){
     .Call(`_arrow_dataset___DSDiscovery__Finish2` , discovery, schema)
 }
 
+dataset___DataSource__schema <- function(source){
+    .Call(`_arrow_dataset___DataSource__schema` , source)
+}
+
 dataset___DSDiscovery__Inspect <- function(discovery){
     .Call(`_arrow_dataset___DSDiscovery__Inspect` , discovery)
 }

r/src/arrowExports.cpp

@@ -1450,6 +1450,21 @@ RcppExport SEXP _arrow_dataset___DSDiscovery__Finish2(SEXP discovery_sexp, SEXP
 }
 #endif
 
+// dataset.cpp
+#if defined(ARROW_R_WITH_ARROW)
+std::shared_ptr<arrow::Schema> dataset___DataSource__schema(const std::shared_ptr<ds::DataSource>& source);
+RcppExport SEXP _arrow_dataset___DataSource__schema(SEXP source_sexp){
+BEGIN_RCPP
+	Rcpp::traits::input_parameter<const std::shared_ptr<ds::DataSource>&>::type source(source_sexp);
+	return Rcpp::wrap(dataset___DataSource__schema(source));
+END_RCPP
+}
+#else
+RcppExport SEXP _arrow_dataset___DataSource__schema(SEXP source_sexp){
+	Rf_error("Cannot call dataset___DataSource__schema(). Please use arrow::install_arrow() to install required runtime libraries. ");
+}
+#endif
+
 // dataset.cpp
 #if defined(ARROW_R_WITH_ARROW)
 std::shared_ptr<arrow::Schema> dataset___DSDiscovery__Inspect(const std::shared_ptr<ds::DataSourceDiscovery>& discovery);
@@ -5638,6 +5653,7 @@ static const R_CallMethodDef CallEntries[] = {
 		{ "_arrow_dataset___IpcFileFormat__Make", (DL_FUNC) &_arrow_dataset___IpcFileFormat__Make, 0}, 
 		{ "_arrow_dataset___DSDiscovery__Finish1", (DL_FUNC) &_arrow_dataset___DSDiscovery__Finish1, 1}, 
 		{ "_arrow_dataset___DSDiscovery__Finish2", (DL_FUNC) &_arrow_dataset___DSDiscovery__Finish2, 2}, 
+		{ "_arrow_dataset___DataSource__schema", (DL_FUNC) &_arrow_dataset___DataSource__schema, 1}, 
 		{ "_arrow_dataset___DSDiscovery__Inspect", (DL_FUNC) &_arrow_dataset___DSDiscovery__Inspect, 1}, 
 		{ "_arrow_dataset___SchemaPartitionScheme", (DL_FUNC) &_arrow_dataset___SchemaPartitionScheme, 1}, 
 		{ "_arrow_dataset___SchemaPartitionScheme__MakeDiscovery", (DL_FUNC) &_arrow_dataset___SchemaPartitionScheme__MakeDiscovery, 1}, 

r/src/dataset.cpp

@@ -80,11 +80,11 @@ std::shared_ptr<ds::DataSource> dataset___DSDiscovery__Finish2(
   return VALUE_OR_STOP(discovery->Finish(schema));
 }
 
-// // [[arrow::export]]
-// std::shared_ptr<arrow::Schema> dataset___DataSource__schema(
-//     const std::shared_ptr<ds::DataSource>& source) {
-//   return source->schema();
-// }
+// [[arrow::export]]
+std::shared_ptr<arrow::Schema> dataset___DataSource__schema(
+    const std::shared_ptr<ds::DataSource>& source) {
+  return source->schema();
+}
 
 // [[arrow::export]]
 std::shared_ptr<arrow::Schema> dataset___DSDiscovery__Inspect(