87 lines
2.8 KiB
C
87 lines
2.8 KiB
C
|
// 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 <string>
|
||
|
|
|
||
|
|
#include "arrow/acero/options.h"
|
||
|
|
#include "arrow/acero/test_util_internal.h"
|
||
|
|
#include "arrow/testing/random.h"
|
||
|
|
|
||
|
|
namespace arrow {
|
||
|
|
namespace acero {
|
||
|
|
|
||
|
|
// \brief Make a delaying source that is optionally noisy (prints when it emits)
|
||
|
|
AsyncGenerator<std::optional<ExecBatch>> MakeDelayedGen(
|
||
|
|
Iterator<std::optional<ExecBatch>> src, std::string label, double delay_sec,
|
||
|
|
bool noisy = false);
|
||
|
|
|
||
|
|
// \brief Make a delaying source that is optionally noisy (prints when it emits)
|
||
|
|
AsyncGenerator<std::optional<ExecBatch>> MakeDelayedGen(
|
||
|
|
AsyncGenerator<std::optional<ExecBatch>> src, std::string label, double delay_sec,
|
||
|
|
bool noisy = false);
|
||
|
|
|
||
|
|
// \brief Make a delaying source that is optionally noisy (prints when it emits)
|
||
|
|
AsyncGenerator<std::optional<ExecBatch>> MakeDelayedGen(BatchesWithSchema src,
|
||
|
|
std::string label,
|
||
|
|
double delay_sec,
|
||
|
|
bool noisy = false);
|
||
|
|
|
||
|
|
/// A node that slightly resequences the input at random
|
||
|
|
struct JitterNodeOptions : public ExecNodeOptions {
|
||
|
|
random::SeedType seed;
|
||
|
|
/// The max amount to add to a node's "cost".
|
||
|
|
int max_jitter_modifier;
|
||
|
|
|
||
|
|
explicit JitterNodeOptions(random::SeedType seed, int max_jitter_modifier = 5)
|
||
|
|
: seed(seed), max_jitter_modifier(max_jitter_modifier) {}
|
||
|
|
static constexpr std::string_view kName = "jitter";
|
||
|
|
};
|
||
|
|
|
||
|
|
class GateImpl;
|
||
|
|
|
||
|
|
class Gate {
|
||
|
|
public:
|
||
|
|
static std::shared_ptr<Gate> Make();
|
||
|
|
|
||
|
|
Gate();
|
||
|
|
virtual ~Gate();
|
||
|
|
|
||
|
|
void ReleaseAllBatches();
|
||
|
|
void ReleaseOneBatch();
|
||
|
|
Future<> WaitForNextReleasedBatch();
|
||
|
|
|
||
|
|
private:
|
||
|
|
ARROW_DISALLOW_COPY_AND_ASSIGN(Gate);
|
||
|
|
|
||
|
|
GateImpl* impl_;
|
||
|
|
};
|
||
|
|
|
||
|
|
// A node that holds all input batches until a given gate is released
|
||
|
|
struct GatedNodeOptions : public ExecNodeOptions {
|
||
|
|
explicit GatedNodeOptions(Gate* gate) : gate(gate) {}
|
||
|
|
Gate* gate;
|
||
|
|
|
||
|
|
static constexpr std::string_view kName = "gated";
|
||
|
|
};
|
||
|
|
|
||
|
|
void RegisterTestNodes();
|
||
|
|
|
||
|
|
} // namespace acero
|
||
|
|
} // namespace arrow
|