arrow/dataset/file_ipc.h

0001 // Licensed to the Apache Software Foundation (ASF) under one
0002 // or more contributor license agreements.  See the NOTICE file
0003 // distributed with this work for additional information
0004 // regarding copyright ownership.  The ASF licenses this file
0005 // to you under the Apache License, Version 2.0 (the
0006 // "License"); you may not use this file except in compliance
0007 // with the License.  You may obtain a copy of the License at
0008 //
0009 //   http://www.apache.org/licenses/LICENSE-2.0
0010 //
0011 // Unless required by applicable law or agreed to in writing,
0012 // software distributed under the License is distributed on an
0013 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
0014 // KIND, either express or implied.  See the License for the
0015 // specific language governing permissions and limitations
0016 // under the License.
0017
0018 // This API is EXPERIMENTAL.
0019
0020 #pragma once
0021
0022 #include <memory>
0023 #include <string>
0024
0025 #include "arrow/dataset/file_base.h"
0026 #include "arrow/dataset/type_fwd.h"
0027 #include "arrow/dataset/visibility.h"
0028 #include "arrow/io/type_fwd.h"
0029 #include "arrow/ipc/type_fwd.h"
0030 #include "arrow/result.h"
0031
0032 namespace arrow {
0033 namespace dataset {
0034
0035 /// \addtogroup dataset-file-formats
0036 ///
0037 /// @{
0038
0039 constexpr char kIpcTypeName[] = "ipc";
0040
0041 /// \brief A FileFormat implementation that reads from and writes to Ipc files
0042 class ARROW_DS_EXPORT IpcFileFormat : public FileFormat {
0043  public:
0044   std::string type_name() const override { return kIpcTypeName; }
0045
0046   IpcFileFormat();
0047
0048   bool Equals(const FileFormat& other) const override {
0049     return type_name() == other.type_name();
0050   }
0051
0052   Result<bool> IsSupported(const FileSource& source) const override;
0053
0054   /// \brief Return the schema of the file if possible.
0055   Result<std::shared_ptr<Schema>> Inspect(const FileSource& source) const override;
0056
0057   Result<RecordBatchGenerator> ScanBatchesAsync(
0058       const std::shared_ptr<ScanOptions>& options,
0059       const std::shared_ptr<FileFragment>& file) const override;
0060
0061   Future<std::optional<int64_t>> CountRows(
0062       const std::shared_ptr<FileFragment>& file, compute::Expression predicate,
0063       const std::shared_ptr<ScanOptions>& options) override;
0064
0065   Result<std::shared_ptr<FileWriter>> MakeWriter(
0066       std::shared_ptr<io::OutputStream> destination, std::shared_ptr<Schema> schema,
0067       std::shared_ptr<FileWriteOptions> options,
0068       fs::FileLocator destination_locator) const override;
0069
0070   std::shared_ptr<FileWriteOptions> DefaultWriteOptions() override;
0071 };
0072
0073 /// \brief Per-scan options for IPC fragments
0074 class ARROW_DS_EXPORT IpcFragmentScanOptions : public FragmentScanOptions {
0075  public:
0076   std::string type_name() const override { return kIpcTypeName; }
0077
0078   /// Options passed to the IPC file reader.
0079   /// included_fields, memory_pool, and use_threads are ignored.
0080   std::shared_ptr<ipc::IpcReadOptions> options;
0081   /// If present, the async scanner will enable I/O coalescing.
0082   /// This is ignored by the sync scanner.
0083   std::shared_ptr<io::CacheOptions> cache_options;
0084 };
0085
0086 class ARROW_DS_EXPORT IpcFileWriteOptions : public FileWriteOptions {
0087  public:
0088   /// Options passed to ipc::MakeFileWriter. use_threads is ignored
0089   std::shared_ptr<ipc::IpcWriteOptions> options;
0090
0091   /// custom_metadata written to the file's footer
0092   std::shared_ptr<const KeyValueMetadata> metadata;
0093
0094  protected:
0095   explicit IpcFileWriteOptions(std::shared_ptr<FileFormat> format)
0096       : FileWriteOptions(std::move(format)) {}
0097
0098   friend class IpcFileFormat;
0099 };
0100
0101 class ARROW_DS_EXPORT IpcFileWriter : public FileWriter {
0102  public:
0103   Status Write(const std::shared_ptr<RecordBatch>& batch) override;
0104
0105  private:
0106   IpcFileWriter(std::shared_ptr<io::OutputStream> destination,
0107                 std::shared_ptr<ipc::RecordBatchWriter> writer,
0108                 std::shared_ptr<Schema> schema,
0109                 std::shared_ptr<IpcFileWriteOptions> options,
0110                 fs::FileLocator destination_locator);
0111
0112   Future<> FinishInternal() override;
0113
0114   std::shared_ptr<io::OutputStream> destination_;
0115   std::shared_ptr<ipc::RecordBatchWriter> batch_writer_;
0116
0117   friend class IpcFileFormat;
0118 };
0119
0120 /// @}
0121
0122 }  // namespace dataset
0123 }  // namespace arrow