use super::{Operator, OperatorResult};
use crate::execution::DataChunk;
use grafeo_common::types::{LogicalType, NodeId};
pub struct SingleRowOperator {
produced: bool,
}
impl SingleRowOperator {
#[must_use]
pub fn new() -> Self {
Self { produced: false }
}
}
impl Default for SingleRowOperator {
fn default() -> Self {
Self::new()
}
}
impl Operator for SingleRowOperator {
fn next(&mut self) -> OperatorResult {
if self.produced {
return Ok(None);
}
self.produced = true;
let mut chunk = DataChunk::with_capacity(&[], 1);
chunk.set_count(1);
Ok(Some(chunk))
}
fn reset(&mut self) {
self.produced = false;
}
fn name(&self) -> &'static str {
"SingleRow"
}
}
pub struct EmptyOperator;
impl EmptyOperator {
#[must_use]
pub fn new(_schema: Vec<LogicalType>) -> Self {
Self
}
}
impl Operator for EmptyOperator {
fn next(&mut self) -> OperatorResult {
Ok(None)
}
fn reset(&mut self) {
}
fn name(&self) -> &'static str {
"Empty"
}
}
pub struct NodeListOperator {
nodes: Vec<NodeId>,
position: usize,
chunk_size: usize,
}
impl NodeListOperator {
#[must_use]
pub fn new(nodes: Vec<NodeId>, chunk_size: usize) -> Self {
Self {
nodes,
position: 0,
chunk_size,
}
}
}
impl Operator for NodeListOperator {
fn next(&mut self) -> OperatorResult {
if self.position >= self.nodes.len() {
return Ok(None);
}
let end = (self.position + self.chunk_size).min(self.nodes.len());
let count = end - self.position;
let schema = [LogicalType::Node];
let mut chunk = DataChunk::with_capacity(&schema, self.chunk_size);
{
let col = chunk
.column_mut(0)
.expect("column 0 exists: chunk created with single-column schema");
for i in self.position..end {
col.push_node_id(self.nodes[i]);
}
}
chunk.set_count(count);
self.position = end;
Ok(Some(chunk))
}
fn reset(&mut self) {
self.position = 0;
}
fn name(&self) -> &'static str {
"NodeList"
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_single_row_operator() {
let mut op = SingleRowOperator::new();
let chunk = op.next().unwrap();
assert!(chunk.is_some());
let chunk = chunk.unwrap();
assert_eq!(chunk.row_count(), 1);
let chunk = op.next().unwrap();
assert!(chunk.is_none());
op.reset();
let chunk = op.next().unwrap();
assert!(chunk.is_some());
}
#[test]
fn test_empty_operator() {
let mut op = EmptyOperator::new(vec![LogicalType::Int64]);
let chunk = op.next().unwrap();
assert!(chunk.is_none());
op.reset();
let chunk = op.next().unwrap();
assert!(chunk.is_none());
}
#[test]
fn test_node_list_operator() {
let nodes = vec![NodeId::new(1), NodeId::new(5), NodeId::new(10)];
let mut op = NodeListOperator::new(nodes, 2);
let chunk = op.next().unwrap();
assert!(chunk.is_some());
let chunk = chunk.unwrap();
assert_eq!(chunk.row_count(), 2);
let chunk = op.next().unwrap();
assert!(chunk.is_some());
let chunk = chunk.unwrap();
assert_eq!(chunk.row_count(), 1);
let chunk = op.next().unwrap();
assert!(chunk.is_none());
op.reset();
let chunk = op.next().unwrap();
assert!(chunk.is_some());
let chunk = chunk.unwrap();
assert_eq!(chunk.row_count(), 2);
}
#[test]
fn test_node_list_operator_empty() {
let mut op = NodeListOperator::new(vec![], 10);
let chunk = op.next().unwrap();
assert!(chunk.is_none());
}
}