#pragma once

#include <functional>

#include <memory>

#include <optional>

#include <string>

#include <vector>

#include "arrow/compute/api_aggregate.h"

#include "arrow/compute/api_vector.h"

#include "arrow/compute/exec.h"

#include "arrow/compute/exec/expression.h"

#include "arrow/record_batch.h"

#include "arrow/result.h"

#include "arrow/util/async_generator.h"

#include "arrow/util/async_util.h"

#include "arrow/util/visibility.h"

namespace arrow {

namespace internal {

class Executor;

} // namespace internal

namespace compute {

using AsyncExecBatchGenerator = AsyncGenerator<std::optional<ExecBatch>>;

class ARROW_EXPORT ExecNodeOptions {

public:

virtual ~ExecNodeOptions() = default;

};

class ARROW_EXPORT SourceNodeOptions : public ExecNodeOptions {

public:

SourceNodeOptions(std::shared_ptr<Schema> output_schema,

std::function<Future<std::optional<ExecBatch>>()> generator)

: output_schema(std::move(output_schema)), generator(std::move(generator)) {}

static Result<std::shared_ptr<SourceNodeOptions>> FromTable(const Table& table,

arrow::internal::Executor*);

static Result<std::shared_ptr<SourceNodeOptions>> FromRecordBatchReader(

std::shared_ptr<RecordBatchReader> reader, std::shared_ptr<Schema> schema,

arrow::internal::Executor*);

std::shared_ptr<Schema> output_schema;

std::function<Future<std::optional<ExecBatch>>()> generator;

};

class ARROW_EXPORT TableSourceNodeOptions : public ExecNodeOptions {

public:

static constexpr int64_t kDefaultMaxBatchSize = 1 << 20;

TableSourceNodeOptions(std::shared_ptr<Table> table,

int64_t max_batch_size = kDefaultMaxBatchSize)

: table(table), max_batch_size(max_batch_size) {}

// arrow table which acts as the data source

std::shared_ptr<Table> table;

// Size of batches to emit from this node

// If the table is larger the node will emit multiple batches from the

// the table to be processed in parallel.

int64_t max_batch_size;

};

class ARROW_EXPORT NamedTableNodeOptions : public ExecNodeOptions {

public:

NamedTableNodeOptions(std::vector<std::string> names, std::shared_ptr<Schema> schema)

: names(std::move(names)), schema(schema) {}

std::vector<std::string> names;

std::shared_ptr<Schema> schema;

};

template <typename ItMaker>

class ARROW_EXPORT SchemaSourceNodeOptions : public ExecNodeOptions {

public:

SchemaSourceNodeOptions(std::shared_ptr<Schema> schema, ItMaker it_maker,

arrow::internal::Executor* io_executor = NULLPTR)

: schema(schema), it_maker(std::move(it_maker)), io_executor(io_executor) {}

/// \brief The schema of the record batches from the iterator

std::shared_ptr<Schema> schema;

/// \brief A maker of an iterator which acts as the data source

ItMaker it_maker;

/// \brief The executor to use for scanning the iterator

///

/// Defaults to the default I/O executor.

arrow::internal::Executor* io_executor;

};

class ARROW_EXPORT RecordBatchReaderSourceNodeOptions : public ExecNodeOptions {

public:

RecordBatchReaderSourceNodeOptions(std::shared_ptr<RecordBatchReader> reader,

arrow::internal::Executor* io_executor = NULLPTR)

: reader(std::move(reader)), io_executor(io_executor) {}

/// \brief The RecordBatchReader which acts as the data source

std::shared_ptr<RecordBatchReader> reader;

/// \brief The executor to use for the reader

///

/// Defaults to the default I/O executor.

arrow::internal::Executor* io_executor;

};

using ArrayVectorIteratorMaker = std::function<Iterator<std::shared_ptr<ArrayVector>>()>;

class ARROW_EXPORT ArrayVectorSourceNodeOptions

: public SchemaSourceNodeOptions<ArrayVectorIteratorMaker> {

using SchemaSourceNodeOptions::SchemaSourceNodeOptions;

};

using ExecBatchIteratorMaker = std::function<Iterator<std::shared_ptr<ExecBatch>>()>;

class ARROW_EXPORT ExecBatchSourceNodeOptions

: public SchemaSourceNodeOptions<ExecBatchIteratorMaker> {

using SchemaSourceNodeOptions::SchemaSourceNodeOptions;

};

using RecordBatchIteratorMaker = std::function<Iterator<std::shared_ptr<RecordBatch>>()>;

class ARROW_EXPORT RecordBatchSourceNodeOptions

: public SchemaSourceNodeOptions<RecordBatchIteratorMaker> {

using SchemaSourceNodeOptions::SchemaSourceNodeOptions;

};

class ARROW_EXPORT FilterNodeOptions : public ExecNodeOptions {

public:

explicit FilterNodeOptions(Expression filter_expression)

: filter_expression(std::move(filter_expression)) {}

Expression filter_expression;

};

class ARROW_EXPORT ProjectNodeOptions : public ExecNodeOptions {

public:

explicit ProjectNodeOptions(std::vector<Expression> expressions,

std::vector<std::string> names = {})

: expressions(std::move(expressions)), names(std::move(names)) {}

std::vector<Expression> expressions;

std::vector<std::string> names;

};

class ARROW_EXPORT AggregateNodeOptions : public ExecNodeOptions {

public:

explicit AggregateNodeOptions(std::vector<Aggregate> aggregates,

std::vector<FieldRef> keys = {})

: aggregates(std::move(aggregates)), keys(std::move(keys)) {}

// aggregations which will be applied to the targetted fields

std::vector<Aggregate> aggregates;

// keys by which aggregations will be grouped

std::vector<FieldRef> keys;

};

constexpr int32_t kDefaultBackpressureHighBytes = 1 << 30; // 1GiB

constexpr int32_t kDefaultBackpressureLowBytes = 1 << 28; // 256MiB

class ARROW_EXPORT BackpressureMonitor {

public:

virtual ~BackpressureMonitor() = default;

virtual uint64_t bytes_in_use() = 0;

virtual bool is_paused() = 0;

};

struct ARROW_EXPORT BackpressureOptions {

/// \brief Create default options that perform no backpressure

BackpressureOptions() : resume_if_below(0), pause_if_above(0) {}

/// \brief Create options that will perform backpressure

///

/// \param resume_if_below The producer should resume producing if the backpressure

/// queue has fewer than resume_if_below items.

/// \param pause_if_above The producer should pause producing if the backpressure

/// queue has more than pause_if_above items

BackpressureOptions(uint64_t resume_if_below, uint64_t pause_if_above)

: resume_if_below(resume_if_below), pause_if_above(pause_if_above) {}

static BackpressureOptions DefaultBackpressure() {

return BackpressureOptions(kDefaultBackpressureLowBytes,

kDefaultBackpressureHighBytes);

}

bool should_apply_backpressure() const { return pause_if_above > 0; }

uint64_t resume_if_below;

uint64_t pause_if_above;

};

class ARROW_EXPORT SinkNodeOptions : public ExecNodeOptions {

public:

explicit SinkNodeOptions(std::function<Future<std::optional<ExecBatch>>()>* generator,

std::shared_ptr<Schema>* schema,

BackpressureOptions backpressure = {},

BackpressureMonitor** backpressure_monitor = NULLPTR)

: generator(generator),

schema(schema),

backpressure(backpressure),

backpressure_monitor(backpressure_monitor) {}

explicit SinkNodeOptions(std::function<Future<std::optional<ExecBatch>>()>* generator,

backpressure_monitor(backpressure_monitor) {}

/// \brief A pointer to a generator of batches.

///

/// This will be set when the node is added to the plan and should be used to consume

/// data from the plan. If this function is not called frequently enough then the sink

/// node will start to accumulate data and may apply backpressure.

std::function<Future<std::optional<ExecBatch>>()>* generator;

/// \brief A pointer which will be set to the schema of the generated batches

///

/// This is optional, if nullptr is passed in then it will be ignored.

/// This will be set when the node is added to the plan, before StartProducing is called

std::shared_ptr<Schema>* schema;

/// \brief Options to control when to apply backpressure

///

/// This is optional, the default is to never apply backpressure. If the plan is not

/// consumed quickly enough the system may eventually run out of memory.

BackpressureOptions backpressure;

/// \brief A pointer to a backpressure monitor

///

/// This will be set when the node is added to the plan. This can be used to inspect

/// the amount of data currently queued in the sink node. This is an optional utility

/// and backpressure can be applied even if this is not used.

BackpressureMonitor** backpressure_monitor;

};

class ARROW_EXPORT BackpressureControl {

public:

virtual ~BackpressureControl() = default;

/// \brief Ask the input to pause

///

/// This is best effort, batches may continue to arrive

/// Must eventually be followed by a call to Resume() or deadlock will occur

virtual void Pause() = 0;

/// \brief Ask the input to resume

virtual void Resume() = 0;

};

class ARROW_EXPORT SinkNodeConsumer {

public:

virtual ~SinkNodeConsumer() = default;

/// \brief Prepare any consumer state

///

/// This will be run once the schema is finalized as the plan is starting and

/// before any calls to Consume. A common use is to save off the schema so that

/// batches can be interpreted.

/// TODO(ARROW-17837) Move ExecPlan* plan to query context

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

BackpressureControl* backpressure_control, ExecPlan* plan) = 0;

/// \brief Consume a batch of data

virtual Status Consume(ExecBatch batch) = 0;

/// \brief Signal to the consumer that the last batch has been delivered

///

/// The returned future should only finish when all outstanding tasks have completed

///

/// If the plan is ended early or aborts due to an error then this will not be

/// called.

virtual Future<> Finish() = 0;

};

class ARROW_EXPORT ConsumingSinkNodeOptions : public ExecNodeOptions {

public:

explicit ConsumingSinkNodeOptions(std::shared_ptr<SinkNodeConsumer> consumer,

std::vector<std::string> names = {})

: consumer(std::move(consumer)), names(std::move(names)) {}

std::shared_ptr<SinkNodeConsumer> consumer;

/// \brief Names to rename the sink's schema fields to

///

/// If specified then names must be provided for all fields. Currently, only a flat

/// schema is supported (see ARROW-15901).

std::vector<std::string> names;

};

class ARROW_EXPORT OrderBySinkNodeOptions : public SinkNodeOptions {

public:

explicit OrderBySinkNodeOptions(

SortOptions sort_options,

std::function<Future<std::optional<ExecBatch>>()>* generator)

: SinkNodeOptions(generator), sort_options(std::move(sort_options)) {}

SortOptions sort_options;

};

enum class JoinType {

};

std::string ToString(JoinType t);

enum class JoinKeyCmp { EQ, IS };

class ARROW_EXPORT HashJoinNodeOptions : public ExecNodeOptions {

public:

static constexpr const char* default_output_suffix_for_left = "";

static constexpr const char* default_output_suffix_for_right = "";

HashJoinNodeOptions(

}

HashJoinNodeOptions(std::vector<FieldRef> in_left_keys,

}

HashJoinNodeOptions(

}

HashJoinNodeOptions(

disable_bloom_filter(disable_bloom_filter) {}

HashJoinNodeOptions() = default;

// type of join (inner, left, semi...)

JoinType join_type = JoinType::INNER;

// key fields from left input

std::vector<FieldRef> left_keys;

// key fields from right input

std::vector<FieldRef> right_keys;

// if set all valid fields from both left and right input will be output

// (and field ref vectors for output fields will be ignored)

bool output_all = false;

// output fields passed from left input

std::vector<FieldRef> left_output;

// output fields passed from right input

std::vector<FieldRef> right_output;

// key comparison function (determines whether a null key is equal another null

// key or not)

std::vector<JoinKeyCmp> key_cmp;

// suffix added to names of output fields coming from left input (used to distinguish,

// if necessary, between fields of the same name in left and right input and can be left

// empty if there are no name collisions)

std::string output_suffix_for_left;

// suffix added to names of output fields coming from right input

std::string output_suffix_for_right;

// residual filter which is applied to matching rows. Rows that do not match

// the filter are not included. The filter is applied against the

// concatenated input schema (left fields then right fields) and can reference

// fields that are not included in the output.

Expression filter = literal(true);

// whether or not to disable Bloom filters in this join

bool disable_bloom_filter = false;

};

class ARROW_EXPORT AsofJoinNodeOptions : public ExecNodeOptions {

public:

/// \brief Keys for one input table of the AsofJoin operation

///

/// The keys must be consistent across the input tables:

/// Each "on" key must refer to a field of the same type and units across the tables.

/// Each "by" key must refer to a list of fields of the same types across the tables.

struct Keys {

};

AsofJoinNodeOptions(std::vector<Keys> input_keys, int64_t tolerance)

: input_keys(std::move(input_keys)), tolerance(tolerance) {}

/// \brief AsofJoin keys per input table. At least two keys must be given. The first key

/// corresponds to a left table and all other keys correspond to right tables for the

/// as-of-join.

///

/// \see `Keys` for details.

std::vector<Keys> input_keys;

/// \brief Tolerance for inexact "on" key matching. A right row is considered a match

/// with the left row if `right.on - left.on <= tolerance`. The `tolerance` may be:

/// - negative, in which case a past-as-of-join occurs;

/// - or positive, in which case a future-as-of-join occurs;

/// - or zero, in which case an exact-as-of-join occurs.

///

/// The tolerance is interpreted in the same units as the "on" key.

int64_t tolerance;

};

class ARROW_EXPORT SelectKSinkNodeOptions : public SinkNodeOptions {

public:

explicit SelectKSinkNodeOptions(

SelectKOptions select_k_options,

std::function<Future<std::optional<ExecBatch>>()>* generator)

: SinkNodeOptions(generator), select_k_options(std::move(select_k_options)) {}

/// SelectK options

SelectKOptions select_k_options;

};

class ARROW_EXPORT TableSinkNodeOptions : public ExecNodeOptions {

public:

explicit TableSinkNodeOptions(std::shared_ptr<Table>* output_table)

: output_table(output_table) {}

std::shared_ptr<Table>* output_table;

};

} // namespace compute

} // namespace arrow