Compression::type compression)

: custom_open_([=] { return ToResult(file); }),

custom_size_(-1),

compression_(compression) {

Result<int64_t> maybe_size = file->GetSize();

if (maybe_size.ok()) {

custom_size_ = *maybe_size;

} else {

custom_open_ = [st = maybe_size.status()] { return st; };

}

if (filesystem_) {

return filesystem_->OpenInputFile(file_info_);

}

if (buffer_) {

return std::make_shared<io::BufferReader>(buffer_);

}

return custom_open_();

if (filesystem_) {

return file_info_.size();

}

if (buffer_) {

return buffer_->size();

}

return custom_size_;

std::optional<Compression::type> compression) const {

ARROW_ASSIGN_OR_RAISE(auto file, Open());

auto actual_compression = Compression::type::UNCOMPRESSED;

if (!compression.has_value()) {

// Guess compression from file extension

auto extension = fs::internal::GetAbstractPathExtension(path());

if (extension == "gz") {

actual_compression = Compression::type::GZIP;

} else {

auto maybe_compression = util::Codec::GetCompressionType(extension);

if (maybe_compression.ok()) {

ARROW_ASSIGN_OR_RAISE(actual_compression, maybe_compression);

}

}

} else {

actual_compression = compression.value();

}

if (actual_compression == Compression::type::UNCOMPRESSED) {

return file;

}

ARROW_ASSIGN_OR_RAISE(auto codec, util::Codec::Create(actual_compression));

return io::CompressedInputStream::Make(codec.get(), std::move(file));

bool match_file_system =

(filesystem_ == nullptr && other.filesystem_ == nullptr) ||

(filesystem_ && other.filesystem_ && filesystem_->Equals(other.filesystem_));

bool match_buffer = (buffer_ == nullptr && other.buffer_ == nullptr) ||

((buffer_ != nullptr && other.buffer_ != nullptr) &&

(buffer_->address() == other.buffer_->address()));

return match_file_system && match_buffer && file_info_.Equals(other.file_info_) &&

compression_ == other.compression_;

const std::shared_ptr<FileFragment>&, compute::Expression,

const std::shared_ptr<ScanOptions>&) {

return Future<std::optional<int64_t>>::MakeFinished(std::nullopt);

const FileSource& source, const FragmentScanOptions* format_options,

compute::ExecContext* exec_context) const {

return Status::NotImplemented("This format does not yet support the scan2 node");

const FragmentScanRequest& request, const InspectedFragment& inspected_fragment,

const FragmentScanOptions* format_options, compute::ExecContext* exec_context) const {

return Status::NotImplemented("This format does not yet support the scan2 node");

FileSource source, std::shared_ptr<Schema> physical_schema) {

return MakeFragment(std::move(source), compute::literal(true),

std::move(physical_schema));

FileSource source, compute::Expression partition_expression,

std::shared_ptr<Schema> physical_schema) {

return std::shared_ptr<FileFragment>(

new FileFragment(std::move(source), shared_from_this(),

std::move(partition_expression), std::move(physical_schema)));

const std::shared_ptr<ScanOptions>& options) {

auto self = std::dynamic_pointer_cast<FileFragment>(shared_from_this());

return format_->ScanBatchesAsync(options, self);

const FragmentScanOptions* format_options, compute::ExecContext* exec_context) {

const FragmentScanOptions* realized_format_options = format_options;

if (format_options == nullptr) {

realized_format_options = format_->default_fragment_scan_options.get();

}

return format_->InspectFragment(source_, realized_format_options, exec_context);

const FragmentScanRequest& request, const InspectedFragment& inspected_fragment,

const FragmentScanOptions* format_options, compute::ExecContext* exec_context) {

const FragmentScanOptions* realized_format_options = format_options;

if (format_options == nullptr) {

realized_format_options = format_->default_fragment_scan_options.get();

}

return format_->BeginScan(request, inspected_fragment, realized_format_options,

exec_context);

compute::Expression predicate, const std::shared_ptr<ScanOptions>& options) {

ARROW_ASSIGN_OR_RAISE(predicate, compute::SimplifyWithGuarantee(std::move(predicate),

partition_expression_));

if (!predicate.IsSatisfiable()) {

return Future<std::optional<int64_t>>::MakeFinished(0);

}

auto self = checked_pointer_cast<FileFragment>(shared_from_this());

return format()->CountRows(self, std::move(predicate), options);

std::shared_ptr<Schema> schema, compute::Expression root_partition,

std::shared_ptr<FileFormat> format, std::shared_ptr<fs::FileSystem> filesystem,

std::vector<std::shared_ptr<FileFragment>> fragments,

std::shared_ptr<Partitioning> partitioning) {

std::shared_ptr<FileSystemDataset> out(

new FileSystemDataset(std::move(schema), std::move(root_partition)));

out->format_ = std::move(format);

out->filesystem_ = std::move(filesystem);

out->fragments_ = std::move(fragments);

out->partitioning_ = std::move(partitioning);

out->SetupSubtreePruning();

return out;

std::shared_ptr<Schema> schema) const {

RETURN_NOT_OK(CheckProjectable(*schema_, *schema));

return Make(std::move(schema), partition_expression_, format_, filesystem_, fragments_);

std::vector<std::string> files;

for (const auto& fragment : fragments_) {

files.push_back(fragment->source().path());

}

return files;

std::string repr = "FileSystemDataset:";

if (fragments_.empty()) {

return repr + " []";

}

for (const auto& fragment : fragments_) {

repr += "\n" + fragment->source().path();

const auto& partition = fragment->partition_expression();

if (partition != compute::literal(true)) {

repr += ": " + partition.ToString();

}

}

return repr;

subtrees_ = std::make_shared<FragmentSubtrees>();

compute::SubtreeImpl impl;

auto encoded = impl.EncodeGuarantees(

[&](int index) { return fragments_[index]->partition_expression(); },

static_cast<int>(fragments_.size()));

std::sort(encoded.begin(), encoded.end(), compute::SubtreeImpl::ByGuarantee());

for (const auto& e : encoded) {

if (e.index) {

subtrees_->fragments_and_subtrees.emplace_back(*e.index);

} else {

subtrees_->fragments_and_subtrees.emplace_back(impl.GetSubtreeExpression(e));

}

}

subtrees_->forest = compute::Forest(static_cast<int>(encoded.size()),

compute::SubtreeImpl::IsAncestor{encoded});

compute::Expression predicate) {

if (predicate == compute::literal(true)) {

// trivial predicate; skip subtree pruning

return MakeVectorIterator(FragmentVector(fragments_.begin(), fragments_.end()));

}

std::vector<int> fragment_indices;

std::vector<compute::Expression> predicates{predicate};

RETURN_NOT_OK(subtrees_->forest.Visit(

[&](compute::Forest::Ref ref) -> Result<bool> {

if (auto fragment_index =

std::get_if<int>(&subtrees_->fragments_and_subtrees[ref.i])) {

fragment_indices.push_back(*fragment_index);

return false;

}

const auto& subtree_expr =

std::get<compute::Expression>(subtrees_->fragments_and_subtrees[ref.i]);

ARROW_ASSIGN_OR_RAISE(auto simplified,

SimplifyWithGuarantee(predicates.back(), subtree_expr));

if (!simplified.IsSatisfiable()) {

return false;

}

predicates.push_back(std::move(simplified));

return true;

},

[&](compute::Forest::Ref ref) { predicates.pop_back(); }));

std::sort(fragment_indices.begin(), fragment_indices.end());

FragmentVector fragments(fragment_indices.size());

std::transform(fragment_indices.begin(), fragment_indices.end(), fragments.begin(),

[this](int i) { return fragments_[i]; });

return MakeVectorIterator(std::move(fragments));

while (true) {

ARROW_ASSIGN_OR_RAISE(auto batch, batches->Next());

if (batch == nullptr) break;

RETURN_NOT_OK(Write(batch));

}

return Status::OK();

return FinishInternal().Then([this]() -> Future<> {

ARROW_ASSIGN_OR_RAISE(bytes_written_, destination_->Tell());

return destination_->CloseAsync();

});

if (bytes_written_.has_value()) {

return bytes_written_.value();

} else {

return Status::Invalid("Cannot retrieve bytes written before calling Finish()");

}

std::shared_ptr<RecordBatch> batch, compute::Expression guarantee,

FileSystemDatasetWriteOptions write_options,

std::function<Status(std::shared_ptr<RecordBatch>, const PartitionPathFormat&)>

write) {

ARROW_ASSIGN_OR_RAISE(auto groups, write_options.partitioning->Partition(batch));

batch.reset(); // drop to hopefully conserve memory

if (write_options.max_partitions <= 0) {

return Status::Invalid("max_partitions must be positive (was ",

write_options.max_partitions, ")");

}

if (groups.batches.size() > static_cast<size_t>(write_options.max_partitions)) {

return Status::Invalid("Fragment would be written into ", groups.batches.size(),

" partitions. This exceeds the maximum of ",

write_options.max_partitions);

}

for (std::size_t index = 0; index < groups.batches.size(); index++) {

auto partition_expression = and_(groups.expressions[index], guarantee);

auto next_batch = groups.batches[index];

PartitionPathFormat destination;

ARROW_ASSIGN_OR_RAISE(destination,

write_options.partitioning->Format(partition_expression));

RETURN_NOT_OK(write(next_batch, destination));

}

return Status::OK();

public:

DatasetWritingSinkNodeConsumer(std::shared_ptr<const KeyValueMetadata> custom_metadata,

FileSystemDatasetWriteOptions write_options)

: custom_metadata_(std::move(custom_metadata)),

write_options_(std::move(write_options)) {}

Status Init(const std::shared_ptr<Schema>& schema,

compute::BackpressureControl* backpressure_control,

compute::ExecPlan* plan) override {

if (custom_metadata_) {

schema_ = schema->WithMetadata(custom_metadata_);

} else {

schema_ = schema;

}

ARROW_ASSIGN_OR_RAISE(

dataset_writer_,

internal::DatasetWriter::Make(

write_options_, plan->query_context()->async_scheduler(),

[backpressure_control] { backpressure_control->Pause(); },

[backpressure_control] { backpressure_control->Resume(); }, [] {}));

return Status::OK();

}

Status Consume(compute::ExecBatch batch) override {

ARROW_ASSIGN_OR_RAISE(std::shared_ptr<RecordBatch> record_batch,

batch.ToRecordBatch(schema_));

return WriteNextBatch(std::move(record_batch), batch.guarantee);

}

Future<> Finish() override {

dataset_writer_->Finish();

// Some write tasks may still be in the queue at this point but that is ok.

return Future<>::MakeFinished();

}

private:

Status WriteNextBatch(std::shared_ptr<RecordBatch> batch,

compute::Expression guarantee) {

return WriteBatch(batch, guarantee, write_options_,

[this](std::shared_ptr<RecordBatch> next_batch,

const PartitionPathFormat& destination) {

dataset_writer_->WriteRecordBatch(std::move(next_batch),

destination.directory,

destination.filename);

return Status::OK();

});

}

std::shared_ptr<const KeyValueMetadata> custom_metadata_;

std::unique_ptr<internal::DatasetWriter> dataset_writer_;

FileSystemDatasetWriteOptions write_options_;

Future<> finished_ = Future<>::Make();

std::shared_ptr<Schema> schema_ = nullptr;

std::shared_ptr<Scanner> scanner) {

auto exprs = scanner->options()->projection.call()->arguments;

auto names = checked_cast<const compute::MakeStructOptions*>(

scanner->options()->projection.call()->options.get())

->field_names;

std::shared_ptr<Dataset> dataset = scanner->dataset();

// The projected_schema is currently used by pyarrow to preserve the custom metadata

// when reading from a single input file.

const auto& custom_metadata = scanner->options()->projected_schema->metadata();

compute::Declaration plan = compute::Declaration::Sequence({

{"scan", ScanNodeOptions{dataset, scanner->options()}},

{"filter", compute::FilterNodeOptions{scanner->options()->filter}},

{"project", compute::ProjectNodeOptions{std::move(exprs), std::move(names)}},

{"write", WriteNodeOptions{write_options, custom_metadata}},

});

return compute::DeclarationToStatus(std::move(plan), scanner->options()->use_threads);

std::vector<compute::ExecNode*> inputs,

const compute::ExecNodeOptions& options) {

if (inputs.size() != 1) {

return Status::Invalid("Write SinkNode requires exactly 1 input, got ",

inputs.size());

}

const WriteNodeOptions write_node_options =

checked_cast<const WriteNodeOptions&>(options);

const std::shared_ptr<const KeyValueMetadata>& custom_metadata =

write_node_options.custom_metadata;

const FileSystemDatasetWriteOptions& write_options = write_node_options.write_options;

if (!write_options.partitioning) {

return Status::Invalid("Must provide partitioning");

}

std::shared_ptr<DatasetWritingSinkNodeConsumer> consumer =

std::make_shared<DatasetWritingSinkNodeConsumer>(custom_metadata, write_options);

ARROW_ASSIGN_OR_RAISE(

auto node,

compute::MakeExecNode("consuming_sink", plan, std::move(inputs),

compute::ConsumingSinkNodeOptions{std::move(consumer)}));

return node;

public:

TeeNode(compute::ExecPlan* plan, std::vector<compute::ExecNode*> inputs,

std::shared_ptr<Schema> output_schema,

FileSystemDatasetWriteOptions write_options)

: MapNode(plan, std::move(inputs), std::move(output_schema)),

write_options_(std::move(write_options)) {}

Status StartProducing() override {

ARROW_ASSIGN_OR_RAISE(

dataset_writer_,

internal::DatasetWriter::Make(

write_options_, plan_->query_context()->async_scheduler(),

[this] { Pause(); }, [this] { Resume(); }, [this] { MapNode::Finish(); }));

return MapNode::StartProducing();

}

static Result<compute::ExecNode*> Make(compute::ExecPlan* plan,

std::vector<compute::ExecNode*> inputs,

const compute::ExecNodeOptions& options) {

RETURN_NOT_OK(ValidateExecNodeInputs(plan, inputs, 1, "TeeNode"));

const WriteNodeOptions write_node_options =

checked_cast<const WriteNodeOptions&>(options);

const FileSystemDatasetWriteOptions& write_options = write_node_options.write_options;

const std::shared_ptr<Schema> schema = inputs[0]->output_schema();

return plan->EmplaceNode<TeeNode>(plan, std::move(inputs), std::move(schema),

std::move(write_options));

}

const char* kind_name() const override { return "TeeNode"; }

void Finish() override { dataset_writer_->Finish(); }

Result<compute::ExecBatch> ProcessBatch(compute::ExecBatch batch) override {

ARROW_ASSIGN_OR_RAISE(std::shared_ptr<RecordBatch> record_batch,

batch.ToRecordBatch(output_schema()));

ARROW_RETURN_NOT_OK(WriteNextBatch(std::move(record_batch), batch.guarantee));

return batch;

}

Status WriteNextBatch(std::shared_ptr<RecordBatch> batch,

compute::Expression guarantee) {

return WriteBatch(batch, guarantee, write_options_,

[this](std::shared_ptr<RecordBatch> next_batch,

const PartitionPathFormat& destination) {

dataset_writer_->WriteRecordBatch(

next_batch, destination.directory, destination.filename);

return Status::OK();

});

}

void Pause() { inputs_[0]->PauseProducing(this, ++backpressure_counter_); }

void Resume() { inputs_[0]->ResumeProducing(this, ++backpressure_counter_); }

protected:

std::string ToStringExtra(int indent = 0) const override {

return "base_dir=" + write_options_.base_dir;

}

private:

std::unique_ptr<internal::DatasetWriter> dataset_writer_;

FileSystemDatasetWriteOptions write_options_;

std::atomic<int32_t> backpressure_counter_ = 0;