1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58
use crate::api::element::Record; use crate::api::function::{Context, Function, KeySelectorFunction}; use crate::functions::schema_base::FunctionSchema; #[derive(Debug)] pub struct SchemaBaseKeySelector { field_types: Vec<u8>, key_field_types: Vec<u8>, columns: Vec<usize>, } impl SchemaBaseKeySelector { pub fn new(columns: Vec<usize>, data_types: &[u8]) -> Self { let key_field_types: Vec<u8> = columns.iter().map(|index| data_types[*index]).collect(); SchemaBaseKeySelector { columns, field_types: data_types.to_vec(), key_field_types, } } } impl FunctionSchema for SchemaBaseKeySelector { fn schema_types(&self) -> Vec<u8> { self.key_field_types.clone() } } impl KeySelectorFunction for SchemaBaseKeySelector { fn open(&mut self, _context: &Context) -> crate::api::Result<()> { Ok(()) } fn get_key(&self, record: &mut Record) -> Record { let mut record_key = Record::with_capacity(record.len()); let mut writer = record_key.as_writer(self.key_field_types.as_slice()); let mut reader = record.as_reader(self.field_types.as_slice()); for index in 0..self.columns.len() { writer .set_bytes_raw(reader.get_bytes_raw(self.columns[index]).unwrap()) .unwrap(); } record_key } fn close(&mut self) -> crate::api::Result<()> { Ok(()) } } impl Function for SchemaBaseKeySelector { fn name(&self) -> &str { "SchemaBaseKeySelector" } }