use std::fmt;
use rudb_common::{Error, Result};
use rudb_pipeline::{Morsel, Progress, Sink, Source, Stream};
use rudb_vector::Chunk;
use crate::buffer::Buffered;
use crate::operator::Operator;
use crate::schema::Schema;
pub(crate) struct Pulled<S: Source> {
source: S,
schema: Schema,
morsel: Option<Morsel>,
done: bool,
}
impl<S: Source> Pulled<S> {
pub(crate) fn new(source: S, schema: Schema) -> Self {
Self { source, schema, morsel: None, done: false }
}
}
impl<S: Source> fmt::Debug for Pulled<S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Pulled").field("source", &self.source).finish_non_exhaustive()
}
}
impl<S: Source> Operator for Pulled<S> {
fn schema(&self) -> &Schema {
&self.schema
}
fn next(&mut self) -> Result<Option<Chunk>> {
while !self.done {
let Some(mut morsel) = self.morsel.take().or_else(|| self.source.morsel()) else {
self.done = true;
break;
};
let mut chunk = Chunk::empty(&[]);
match self.source.read(&mut morsel, &mut chunk)? {
Progress::Blocked(blocked) => return Err(parked(&blocked)),
Progress::Done => {}
_ => self.morsel = Some(morsel),
}
if !chunk.is_empty() {
return Ok(Some(chunk));
}
}
Ok(None)
}
}
pub(crate) struct Streamed<'a, S: Stream> {
input: Box<dyn Operator + 'a>,
stream: S,
local: S::Local,
schema: Schema,
done: bool,
again: bool,
}
impl<'a, S: Stream> Streamed<'a, S> {
pub(crate) fn new(input: Box<dyn Operator + 'a>, stream: S, schema: Schema) -> Self {
let local = stream.local();
Self { input, stream, local, schema, done: false, again: false }
}
}
impl<S: Stream> fmt::Debug for Streamed<'_, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Streamed")
.field("stream", &self.stream)
.field("input", &self.input)
.finish_non_exhaustive()
}
}
impl<S: Stream> Operator for Streamed<'_, S> {
fn schema(&self) -> &Schema {
&self.schema
}
fn next(&mut self) -> Result<Option<Chunk>> {
while !self.done {
let mut chunk = if self.again {
Chunk::empty(&[])
} else {
let Some(chunk) = self.input.next()? else { break };
chunk
};
self.again = false;
match self.stream.push(&mut chunk, &mut self.local)? {
Progress::Done => self.done = true,
Progress::Again => self.again = true,
Progress::Blocked(blocked) => return Err(parked(&blocked)),
_ => {}
}
if !chunk.is_empty() {
return Ok(Some(chunk));
}
}
Ok(None)
}
}
pub(crate) struct Fed<'a, F: Sink, S: Stream> {
first: Box<dyn Operator + 'a>,
aside: F,
then: Streamed<'a, S>,
built: bool,
}
impl<'a, F: Sink, S: Stream> Fed<'a, F, S> {
pub(crate) fn new(first: Box<dyn Operator + 'a>, aside: F, then: Streamed<'a, S>) -> Self {
Self { first, aside, then, built: false }
}
}
impl<F: Sink, S: Stream> fmt::Debug for Fed<'_, F, S> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Fed")
.field("aside", &self.aside)
.field("then", &self.then)
.finish_non_exhaustive()
}
}
impl<F: Sink, S: Stream> Operator for Fed<'_, F, S> {
fn schema(&self) -> &Schema {
self.then.schema()
}
fn next(&mut self) -> Result<Option<Chunk>> {
if !self.built {
drain(self.first.as_mut(), &self.aside)?;
self.built = true;
}
self.then.next()
}
}
pub(crate) struct Broken<'a, K: Sink> {
input: Box<dyn Operator + 'a>,
sink: K,
out: Buffered,
schema: Schema,
built: bool,
at: usize,
}
impl<'a, K: Sink> Broken<'a, K> {
pub(crate) fn new(
input: Box<dyn Operator + 'a>,
sink: K,
out: Buffered,
schema: Schema,
) -> Self {
Self { input, sink, out, schema, built: false, at: 0 }
}
fn build(&mut self) -> Result<()> {
drain(self.input.as_mut(), &self.sink)
}
}
impl<K: Sink> fmt::Debug for Broken<'_, K> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Broken")
.field("sink", &self.sink)
.field("input", &self.input)
.finish_non_exhaustive()
}
}
impl<K: Sink> Operator for Broken<'_, K> {
fn schema(&self) -> &Schema {
&self.schema
}
fn next(&mut self) -> Result<Option<Chunk>> {
if !self.built {
self.build()?;
self.built = true;
}
let chunk = self.out.at(self.at)?;
if chunk.is_some() {
self.at += 1;
}
Ok(chunk)
}
}
pub(crate) struct Paired<'a, F: Sink, K: Sink> {
first: Box<dyn Operator + 'a>,
aside: F,
second: Box<dyn Operator + 'a>,
sink: K,
out: Buffered,
schema: Schema,
built: bool,
at: usize,
}
impl<'a, F: Sink, K: Sink> Paired<'a, F, K> {
pub(crate) fn new(
first: Box<dyn Operator + 'a>,
aside: F,
second: Box<dyn Operator + 'a>,
sink: K,
out: Buffered,
schema: Schema,
) -> Self {
Self { first, aside, second, sink, out, schema, built: false, at: 0 }
}
fn build(&mut self) -> Result<()> {
drain(self.first.as_mut(), &self.aside)?;
drain(self.second.as_mut(), &self.sink)
}
}
impl<F: Sink, K: Sink> fmt::Debug for Paired<'_, F, K> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Paired")
.field("sink", &self.sink)
.field("aside", &self.aside)
.finish_non_exhaustive()
}
}
impl<F: Sink, K: Sink> Operator for Paired<'_, F, K> {
fn schema(&self) -> &Schema {
&self.schema
}
fn next(&mut self) -> Result<Option<Chunk>> {
if !self.built {
self.build()?;
self.built = true;
}
let chunk = self.out.at(self.at)?;
if chunk.is_some() {
self.at += 1;
}
Ok(chunk)
}
}
fn drain<K: Sink>(input: &mut dyn Operator, sink: &K) -> Result<()> {
let mut local = sink.local();
while let Some(chunk) = input.next()? {
match sink.sink(&chunk, &mut local)? {
Progress::Done => break,
Progress::Blocked(blocked) => return Err(parked(&blocked)),
_ => {}
}
}
sink.combine(local)?;
sink.finalize()
}
fn parked(blocked: &rudb_pipeline::Blocked) -> Error {
Error::not_implemented(format!(
"an operator blocked {blocked} and the pull tree has nothing else to run, which is F4"
))
}