use crate::io::FileIO;
use crate::spec::{batch_to_serialized_bytes, DataField, IndexFileMeta, EMPTY_SERIALIZED_ROW};
use crate::table::bucket_assigner::{BatchAssignOutput, BucketAssigner, PartitionBucketKey};
use crate::Result;
use arrow_array::RecordBatch;
use std::collections::HashMap;
pub(crate) struct ConstantBucketAssigner {
partition_field_indices: Vec<usize>,
bucket: i32,
}
impl ConstantBucketAssigner {
pub fn new(partition_field_indices: Vec<usize>, bucket: i32) -> Self {
Self {
partition_field_indices,
bucket,
}
}
pub fn bucket(&self) -> i32 {
self.bucket
}
}
impl BucketAssigner for ConstantBucketAssigner {
async fn assign_batch(
&mut self,
batch: &RecordBatch,
fields: &[DataField],
) -> Result<BatchAssignOutput> {
let num_rows = batch.num_rows();
let partition_bytes = if self.partition_field_indices.is_empty() {
vec![EMPTY_SERIALIZED_ROW.clone(); num_rows]
} else {
batch_to_serialized_bytes(batch, &self.partition_field_indices, fields)?
};
Ok(BatchAssignOutput {
partition_bytes,
buckets: vec![self.bucket; num_rows],
deletes: Vec::new(),
skips: Vec::new(),
})
}
async fn prepare_commit_index(
&mut self,
_file_io: &FileIO,
_index_dir: &str,
) -> Result<HashMap<PartitionBucketKey, Vec<IndexFileMeta>>> {
Ok(HashMap::new())
}
}