use crate::types::{AggOp, Schema};
use crate::Result;
#[derive(Debug)]
pub struct GpuBuffer {
pub num_rows: u64,
pub schema: Schema,
}
impl GpuBuffer {
pub fn empty() -> Self {
Self {
num_rows: 0,
schema: Schema::new(vec![]),
}
}
pub fn is_empty(&self) -> bool {
self.num_rows == 0
}
pub fn estimated_bytes(&self) -> u64 {
self.num_rows * self.schema.row_size_bytes() as u64
}
}
pub trait KernelProvider: Send + Sync {
fn hash_join(
&self,
left: &GpuBuffer,
right: &GpuBuffer,
left_keys: &[usize],
right_keys: &[usize],
) -> Result<GpuBuffer>;
fn dedup(&self, input: &GpuBuffer, key_cols: &[usize]) -> Result<GpuBuffer>;
fn union(&self, a: &GpuBuffer, b: &GpuBuffer) -> Result<GpuBuffer>;
fn diff(&self, a: &GpuBuffer, b: &GpuBuffer) -> Result<GpuBuffer>;
fn groupby_agg(
&self,
input: &GpuBuffer,
key_cols: &[usize],
agg: AggOp,
value_col: usize,
) -> Result<GpuBuffer>;
}
pub trait RelationStore: Send + Sync {
fn get(&self, name: &str) -> Option<&GpuBuffer>;
fn put(&mut self, name: &str, buffer: GpuBuffer);
fn contains(&self, name: &str) -> bool;
}
#[cfg(test)]
mod tests {
use super::*;
struct MockProvider;
impl KernelProvider for MockProvider {
fn hash_join(
&self,
_left: &GpuBuffer,
_right: &GpuBuffer,
_left_keys: &[usize],
_right_keys: &[usize],
) -> Result<GpuBuffer> {
Ok(GpuBuffer::empty())
}
fn dedup(&self, _input: &GpuBuffer, _key_cols: &[usize]) -> Result<GpuBuffer> {
Ok(GpuBuffer::empty())
}
fn union(&self, _a: &GpuBuffer, _b: &GpuBuffer) -> Result<GpuBuffer> {
Ok(GpuBuffer::empty())
}
fn diff(&self, _a: &GpuBuffer, _b: &GpuBuffer) -> Result<GpuBuffer> {
Ok(GpuBuffer::empty())
}
fn groupby_agg(
&self,
_input: &GpuBuffer,
_key_cols: &[usize],
_agg: AggOp,
_value_col: usize,
) -> Result<GpuBuffer> {
Ok(GpuBuffer::empty())
}
}
#[test]
fn test_mock_provider_compiles() {
let provider = MockProvider;
let empty = GpuBuffer::empty();
assert!(provider.dedup(&empty, &[0]).is_ok());
}
#[test]
fn test_gpu_buffer_empty() {
let buf = GpuBuffer::empty();
assert!(buf.is_empty());
assert_eq!(buf.estimated_bytes(), 0);
}
}