use std::sync::Arc;
use thiserror::Error;
#[derive(Clone, Debug, PartialEq, Eq, Error)]
pub enum AttributeError {
#[error("input buffer initial and max must be nonzero")]
ZeroInputBuffer,
#[error("input buffer initial ({initial}) must be <= max ({max})")]
InputBufferMinExceedsMax { initial: usize, max: usize },
#[error("batch size must be greater than zero")]
ZeroBatchSize,
#[error("parallelism must be greater than zero")]
ZeroParallelism,
}
pub type AttributeResult<T> = Result<T, AttributeError>;
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum Fusion {
#[default]
Auto,
TypedOnly,
ErasedOnly,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum MetricsLevel {
#[default]
Off,
Counts,
Timing,
Stalls,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
pub enum KernelBackend {
Auto,
#[default]
ArrowRust,
Python,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Attribute {
Name(Arc<str>),
InputBuffer { initial: usize, max: usize },
Dispatcher(Arc<str>),
BatchSize(usize),
Parallelism(usize),
Fusion(Fusion),
MetricsLevel(MetricsLevel),
KernelBackend(KernelBackend),
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct Attributes {
attribute_list: Arc<[Attribute]>,
}
impl Attributes {
#[must_use]
pub fn none() -> Self {
Self::default()
}
#[must_use]
pub fn named(name: impl Into<String>) -> Self {
Self {
attribute_list: Arc::from([Attribute::Name(Arc::from(name.into()))]),
}
}
#[must_use]
pub fn input_buffer(initial: usize, max: usize) -> Self {
Self::try_input_buffer(initial, max).expect("invalid input buffer attribute")
}
pub fn try_input_buffer(initial: usize, max: usize) -> AttributeResult<Self> {
validate_input_buffer(initial, max)?;
Ok(Self {
attribute_list: Arc::from([Attribute::InputBuffer { initial, max }]),
})
}
#[must_use]
pub fn dispatcher(name: impl Into<String>) -> Self {
Self {
attribute_list: Arc::from([Attribute::Dispatcher(Arc::from(name.into()))]),
}
}
#[must_use]
pub fn batch_size(size: usize) -> Self {
Self::try_batch_size(size).expect("invalid batch size attribute")
}
pub fn try_batch_size(size: usize) -> AttributeResult<Self> {
if size == 0 {
return Err(AttributeError::ZeroBatchSize);
}
Ok(Self {
attribute_list: Arc::from([Attribute::BatchSize(size)]),
})
}
#[must_use]
pub fn parallelism(parallelism: usize) -> Self {
Self::try_parallelism(parallelism).expect("invalid parallelism attribute")
}
pub fn try_parallelism(parallelism: usize) -> AttributeResult<Self> {
if parallelism == 0 {
return Err(AttributeError::ZeroParallelism);
}
Ok(Self {
attribute_list: Arc::from([Attribute::Parallelism(parallelism)]),
})
}
#[must_use]
pub fn fusion(fusion: Fusion) -> Self {
Self {
attribute_list: Arc::from([Attribute::Fusion(fusion)]),
}
}
#[must_use]
pub fn metrics_level(level: MetricsLevel) -> Self {
Self {
attribute_list: Arc::from([Attribute::MetricsLevel(level)]),
}
}
#[must_use]
pub fn kernel_backend(backend: KernelBackend) -> Self {
Self {
attribute_list: Arc::from([Attribute::KernelBackend(backend)]),
}
}
#[must_use]
pub fn and(mut self, other: Self) -> Self {
if other.attribute_list.is_empty() {
return self;
}
if self.attribute_list.is_empty() {
return other;
}
let mut combined =
Vec::with_capacity(other.attribute_list.len() + self.attribute_list.len());
combined.extend(other.attribute_list.iter().cloned());
combined.extend(self.attribute_list.iter().cloned());
self.attribute_list = Arc::from(combined.into_boxed_slice());
self
}
#[must_use]
pub fn attribute_list(&self) -> &[Attribute] {
&self.attribute_list
}
pub fn validate(&self) -> AttributeResult<()> {
for attribute in self.attribute_list.iter() {
match attribute {
Attribute::InputBuffer { initial, max } => validate_input_buffer(*initial, *max)?,
Attribute::BatchSize(0) => return Err(AttributeError::ZeroBatchSize),
Attribute::Parallelism(0) => return Err(AttributeError::ZeroParallelism),
Attribute::Name(_)
| Attribute::Dispatcher(_)
| Attribute::BatchSize(_)
| Attribute::Parallelism(_)
| Attribute::Fusion(_)
| Attribute::MetricsLevel(_)
| Attribute::KernelBackend(_) => {}
}
}
Ok(())
}
#[must_use]
pub fn name(&self) -> Option<&str> {
self.attribute_list
.iter()
.find_map(|attribute| match attribute {
Attribute::Name(name) => Some(name.as_ref()),
_ => None,
})
}
#[must_use]
pub fn input_buffer_hint(&self) -> Option<(usize, usize)> {
self.attribute_list
.iter()
.find_map(|attribute| match attribute {
Attribute::InputBuffer { initial, max } => Some((*initial, *max)),
_ => None,
})
}
#[must_use]
pub fn dispatcher_hint(&self) -> Option<&str> {
self.attribute_list
.iter()
.find_map(|attribute| match attribute {
Attribute::Dispatcher(name) => Some(name.as_ref()),
_ => None,
})
}
#[must_use]
pub fn batch_size_hint(&self) -> Option<usize> {
self.attribute_list
.iter()
.find_map(|attribute| match attribute {
Attribute::BatchSize(size) => Some(*size),
_ => None,
})
}
#[must_use]
pub fn parallelism_hint(&self) -> Option<usize> {
self.attribute_list
.iter()
.find_map(|attribute| match attribute {
Attribute::Parallelism(parallelism) => Some(*parallelism),
_ => None,
})
}
#[must_use]
pub fn fusion_hint(&self) -> Option<Fusion> {
self.attribute_list
.iter()
.find_map(|attribute| match attribute {
Attribute::Fusion(fusion) => Some(*fusion),
_ => None,
})
}
#[must_use]
pub fn metrics_level_hint(&self) -> Option<MetricsLevel> {
self.attribute_list
.iter()
.find_map(|attribute| match attribute {
Attribute::MetricsLevel(level) => Some(*level),
_ => None,
})
}
#[must_use]
pub fn kernel_backend_hint(&self) -> Option<KernelBackend> {
self.attribute_list
.iter()
.find_map(|attribute| match attribute {
Attribute::KernelBackend(backend) => Some(*backend),
_ => None,
})
}
#[must_use]
pub fn with_name(mut self, name: impl Into<String>) -> Self {
let mut attributes = Vec::with_capacity(self.attribute_list.len() + 1);
attributes.push(Attribute::Name(Arc::from(name.into())));
attributes.extend(self.attribute_list.iter().cloned());
self.attribute_list = Arc::from(attributes.into_boxed_slice());
self
}
}
fn validate_input_buffer(initial: usize, max: usize) -> AttributeResult<()> {
if initial == 0 || max == 0 {
return Err(AttributeError::ZeroInputBuffer);
}
if initial > max {
return Err(AttributeError::InputBufferMinExceedsMax { initial, max });
}
Ok(())
}