1use crate::builder::{ArrayBuilder, FixedSizeBinaryBuilder, PrimitiveBuilder};
6use crate::types::ArrowDictionaryKeyType;
7use crate::{Array, ArrayRef, DictionaryArray, PrimitiveArray};
8use arrow_buffer::ArrowNativeType;
9use arrow_schema::DataType::FixedSizeBinary;
10use arrow_schema::{ArrowError, DataType};
11use hashbrown::HashTable;
12use num_traits::NumCast;
13use std::any::Any;
14use std::sync::Arc;
15
16#[derive(Debug)]
48pub struct FixedSizeBinaryDictionaryBuilder<K>
49where
50 K: ArrowDictionaryKeyType,
51{
52 state: ahash::RandomState,
53 dedup: HashTable<usize>,
54
55 keys_builder: PrimitiveBuilder<K>,
56 values_builder: FixedSizeBinaryBuilder,
57 byte_width: i32,
58}
59
60impl<K> FixedSizeBinaryDictionaryBuilder<K>
61where
62 K: ArrowDictionaryKeyType,
63{
64 pub fn new(byte_width: i32) -> Self {
66 let keys_builder = PrimitiveBuilder::new();
67 let values_builder = FixedSizeBinaryBuilder::new(byte_width);
68 Self {
69 state: Default::default(),
70 dedup: HashTable::with_capacity(keys_builder.capacity()),
71 keys_builder,
72 values_builder,
73 byte_width,
74 }
75 }
76
77 pub fn with_capacity(keys_capacity: usize, value_capacity: usize, byte_width: i32) -> Self {
83 Self {
84 state: Default::default(),
85 dedup: Default::default(),
86 keys_builder: PrimitiveBuilder::with_capacity(keys_capacity),
87 values_builder: FixedSizeBinaryBuilder::with_capacity(value_capacity, byte_width),
88 byte_width,
89 }
90 }
91
92 pub fn try_new_from_builder<K2>(
118 mut source: FixedSizeBinaryDictionaryBuilder<K2>,
119 ) -> Result<Self, ArrowError>
120 where
121 K::Native: NumCast,
122 K2: ArrowDictionaryKeyType,
123 K2::Native: NumCast,
124 {
125 let state = source.state;
126 let dedup = source.dedup;
127 let values_builder = source.values_builder;
128 let byte_width = source.byte_width;
129
130 let source_keys = source.keys_builder.finish();
131 let new_keys: PrimitiveArray<K> = source_keys.try_unary(|value| {
132 num_traits::cast::cast::<K2::Native, K::Native>(value).ok_or_else(|| {
133 ArrowError::CastError(format!(
134 "Can't cast dictionary keys from source type {:?} to type {:?}",
135 K2::DATA_TYPE,
136 K::DATA_TYPE
137 ))
138 })
139 })?;
140
141 drop(source_keys);
145
146 Ok(Self {
147 state,
148 dedup,
149 keys_builder: new_keys.into_builder().map_err(|_| {
150 ArrowError::ComputeError(
151 "Internal Error: the keys just derived from the source builder are \
152 unexpectedly shared, so they cannot be reused as a builder"
153 .to_string(),
154 )
155 })?,
156 values_builder,
157 byte_width,
158 })
159 }
160}
161
162impl<K> ArrayBuilder for FixedSizeBinaryDictionaryBuilder<K>
163where
164 K: ArrowDictionaryKeyType,
165{
166 fn as_any(&self) -> &dyn Any {
168 self
169 }
170
171 fn as_any_mut(&mut self) -> &mut dyn Any {
173 self
174 }
175
176 fn into_box_any(self: Box<Self>) -> Box<dyn Any> {
178 self
179 }
180
181 fn len(&self) -> usize {
183 self.keys_builder.len()
184 }
185
186 fn finish(&mut self) -> ArrayRef {
188 Arc::new(self.finish())
189 }
190
191 fn finish_cloned(&self) -> ArrayRef {
193 Arc::new(self.finish_cloned())
194 }
195
196 fn finish_preserve_values(&mut self) -> ArrayRef {
197 Arc::new(self.finish_preserve_values())
198 }
199}
200
201impl<K> FixedSizeBinaryDictionaryBuilder<K>
202where
203 K: ArrowDictionaryKeyType,
204{
205 fn get_or_insert_key(&mut self, value: impl AsRef<[u8]>) -> Result<K::Native, ArrowError> {
206 let value_bytes: &[u8] = value.as_ref();
207
208 let state = &self.state;
209 let storage = &mut self.values_builder;
210 let hash = state.hash_one(value_bytes);
211
212 let idx = *self
213 .dedup
214 .entry(
215 hash,
216 |idx| value_bytes == get_bytes(storage, self.byte_width, *idx),
217 |idx| state.hash_one(get_bytes(storage, self.byte_width, *idx)),
218 )
219 .or_insert_with(|| {
220 let idx = storage.len();
221 let _ = storage.append_value(value);
222 idx
223 })
224 .get();
225
226 let key = K::Native::from_usize(idx).ok_or(ArrowError::DictionaryKeyOverflowError)?;
227
228 Ok(key)
229 }
230
231 pub fn append(&mut self, value: impl AsRef<[u8]>) -> Result<K::Native, ArrowError> {
237 if self.byte_width != value.as_ref().len() as i32 {
238 Err(ArrowError::InvalidArgumentError(format!(
239 "Invalid input length passed to FixedSizeBinaryBuilder. Expected {} got {}",
240 self.byte_width,
241 value.as_ref().len()
242 )))
243 } else {
244 let key = self.get_or_insert_key(value)?;
245 self.keys_builder.append_value(key);
246 Ok(key)
247 }
248 }
249
250 pub fn append_n(
255 &mut self,
256 value: impl AsRef<[u8]>,
257 count: usize,
258 ) -> Result<K::Native, ArrowError> {
259 if self.byte_width != value.as_ref().len() as i32 {
260 Err(ArrowError::InvalidArgumentError(format!(
261 "Invalid input length passed to FixedSizeBinaryBuilder. Expected {} got {}",
262 self.byte_width,
263 value.as_ref().len()
264 )))
265 } else {
266 let key = self.get_or_insert_key(value)?;
267 self.keys_builder.append_value_n(key, count);
268 Ok(key)
269 }
270 }
271
272 #[inline]
274 pub fn append_null(&mut self) {
275 self.keys_builder.append_null()
276 }
277
278 #[inline]
280 pub fn append_nulls(&mut self, n: usize) {
281 self.keys_builder.append_nulls(n);
282 }
283
284 pub fn append_value(&mut self, value: impl AsRef<[u8]>) {
290 self.append(value).expect("dictionary key overflow");
291 }
292
293 pub fn finish(&mut self) -> DictionaryArray<K> {
295 self.dedup.clear();
296 let values = self.values_builder.finish();
297 let keys = self.keys_builder.finish();
298
299 let data_type = DataType::Dictionary(
300 Box::new(K::DATA_TYPE),
301 Box::new(FixedSizeBinary(self.byte_width)),
302 );
303
304 let builder = keys
305 .into_data()
306 .into_builder()
307 .data_type(data_type)
308 .child_data(vec![values.into_data()]);
309
310 DictionaryArray::from(unsafe { builder.build_unchecked() })
312 }
313
314 pub fn finish_cloned(&self) -> DictionaryArray<K> {
316 let values = self.values_builder.finish_cloned();
317 let keys = self.keys_builder.finish_cloned();
318
319 let data_type = DataType::Dictionary(
320 Box::new(K::DATA_TYPE),
321 Box::new(FixedSizeBinary(self.byte_width)),
322 );
323
324 let builder = keys
325 .into_data()
326 .into_builder()
327 .data_type(data_type)
328 .child_data(vec![values.into_data()]);
329
330 DictionaryArray::from(unsafe { builder.build_unchecked() })
332 }
333
334 pub fn finish_preserve_values(&mut self) -> DictionaryArray<K> {
352 let values = self.values_builder.finish_cloned();
353 let keys = self.keys_builder.finish();
354
355 let data_type = DataType::Dictionary(
356 Box::new(K::DATA_TYPE),
357 Box::new(FixedSizeBinary(self.byte_width)),
358 );
359
360 let builder = keys
361 .into_data()
362 .into_builder()
363 .data_type(data_type)
364 .child_data(vec![values.into_data()]);
365
366 DictionaryArray::from(unsafe { builder.build_unchecked() })
368 }
369}
370
371fn get_bytes(values: &FixedSizeBinaryBuilder, byte_width: i32, idx: usize) -> &[u8] {
372 let values = values.values_slice();
373 let start = idx * byte_width.as_usize();
374 let end = idx * byte_width.as_usize() + byte_width.as_usize();
375 &values[start..end]
376}
377
378#[cfg(test)]
379mod tests {
380 use super::*;
381
382 use crate::types::{Int8Type, Int16Type, Int32Type, UInt8Type, UInt16Type};
383 use crate::{ArrowPrimitiveType, FixedSizeBinaryArray, Int8Array};
384
385 #[test]
386 fn test_fixed_size_dictionary_builder() {
387 let values = ["abc", "def"];
388
389 let mut b = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
390 assert_eq!(b.append(values[0]).unwrap(), 0);
391 b.append_null();
392 assert_eq!(b.append(values[1]).unwrap(), 1);
393 assert_eq!(b.append(values[1]).unwrap(), 1);
394 assert_eq!(b.append(values[0]).unwrap(), 0);
395 b.append_nulls(2);
396 assert_eq!(b.append(values[0]).unwrap(), 0);
397 let array = b.finish();
398
399 assert_eq!(
400 array.keys(),
401 &Int8Array::from(vec![
402 Some(0),
403 None,
404 Some(1),
405 Some(1),
406 Some(0),
407 None,
408 None,
409 Some(0)
410 ]),
411 );
412
413 let ava = array
415 .values()
416 .as_any()
417 .downcast_ref::<FixedSizeBinaryArray>()
418 .unwrap();
419
420 assert_eq!(ava.value(0), values[0].as_bytes());
421 assert_eq!(ava.value(1), values[1].as_bytes());
422 }
423
424 #[test]
425 fn test_fixed_size_dictionary_builder_append_n() {
426 let values = ["abc", "def"];
427 let mut b = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
428 assert_eq!(b.append_n(values[0], 2).unwrap(), 0);
429 assert_eq!(b.append_n(values[1], 3).unwrap(), 1);
430 assert_eq!(b.append_n(values[0], 2).unwrap(), 0);
431 let array = b.finish();
432
433 assert_eq!(
434 array.keys(),
435 &Int8Array::from(vec![
436 Some(0),
437 Some(0),
438 Some(1),
439 Some(1),
440 Some(1),
441 Some(0),
442 Some(0),
443 ]),
444 );
445
446 let ava = array
448 .values()
449 .as_any()
450 .downcast_ref::<FixedSizeBinaryArray>()
451 .unwrap();
452
453 assert_eq!(ava.value(0), values[0].as_bytes());
454 assert_eq!(ava.value(1), values[1].as_bytes());
455 }
456
457 #[test]
458 fn test_fixed_size_dictionary_builder_wrong_size() {
459 let mut b = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
460 let err = b.append(b"too long").unwrap_err().to_string();
461 assert_eq!(
462 err,
463 "Invalid argument error: Invalid input length passed to FixedSizeBinaryBuilder. Expected 3 got 8"
464 );
465 let err = b.append("").unwrap_err().to_string();
466 assert_eq!(
467 err,
468 "Invalid argument error: Invalid input length passed to FixedSizeBinaryBuilder. Expected 3 got 0"
469 );
470 let err = b.append_n("a", 3).unwrap_err().to_string();
471 assert_eq!(
472 err,
473 "Invalid argument error: Invalid input length passed to FixedSizeBinaryBuilder. Expected 3 got 1"
474 );
475 }
476
477 #[test]
478 fn test_fixed_size_dictionary_builder_finish_cloned() {
479 let values = ["abc", "def", "ghi"];
480
481 let mut builder = FixedSizeBinaryDictionaryBuilder::<Int8Type>::new(3);
482
483 builder.append(values[0]).unwrap();
484 builder.append_null();
485 builder.append(values[1]).unwrap();
486 builder.append(values[1]).unwrap();
487 builder.append(values[0]).unwrap();
488 let mut array = builder.finish_cloned();
489
490 assert_eq!(
491 array.keys(),
492 &Int8Array::from(vec![Some(0), None, Some(1), Some(1), Some(0)])
493 );
494
495 let ava = array
497 .values()
498 .as_any()
499 .downcast_ref::<FixedSizeBinaryArray>()
500 .unwrap();
501
502 assert_eq!(ava.value(0), values[0].as_bytes());
503 assert_eq!(ava.value(1), values[1].as_bytes());
504
505 builder.append(values[0]).unwrap();
506 builder.append(values[2]).unwrap();
507 builder.append(values[1]).unwrap();
508
509 array = builder.finish();
510
511 assert_eq!(
512 array.keys(),
513 &Int8Array::from(vec![
514 Some(0),
515 None,
516 Some(1),
517 Some(1),
518 Some(0),
519 Some(0),
520 Some(2),
521 Some(1)
522 ])
523 );
524
525 let ava2 = array
527 .values()
528 .as_any()
529 .downcast_ref::<FixedSizeBinaryArray>()
530 .unwrap();
531
532 assert_eq!(ava2.value(0), values[0].as_bytes());
533 assert_eq!(ava2.value(1), values[1].as_bytes());
534 assert_eq!(ava2.value(2), values[2].as_bytes());
535 }
536
537 fn _test_try_new_from_builder_generic_for_key_types<K1, K2>(values: Vec<[u8; 3]>)
538 where
539 K1: ArrowDictionaryKeyType,
540 K1::Native: NumCast,
541 K2: ArrowDictionaryKeyType,
542 K2::Native: NumCast + From<u8>,
543 {
544 let mut source = FixedSizeBinaryDictionaryBuilder::<K1>::new(3);
545 source.append_value(values[0]);
546 source.append_null();
547 source.append_value(values[1]);
548 source.append_value(values[2]);
549
550 let mut result =
551 FixedSizeBinaryDictionaryBuilder::<K2>::try_new_from_builder(source).unwrap();
552 let array = result.finish();
553
554 let mut expected_keys_builder = PrimitiveBuilder::<K2>::new();
555 expected_keys_builder
556 .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(0u8));
557 expected_keys_builder.append_null();
558 expected_keys_builder
559 .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(1u8));
560 expected_keys_builder
561 .append_value(<<K2 as ArrowPrimitiveType>::Native as From<u8>>::from(2u8));
562 let expected_keys = expected_keys_builder.finish();
563 assert_eq!(array.keys(), &expected_keys);
564
565 let av = array.values();
566 let ava = av.as_any().downcast_ref::<FixedSizeBinaryArray>().unwrap();
567 assert_eq!(ava.value(0), values[0]);
568 assert_eq!(ava.value(1), values[1]);
569 assert_eq!(ava.value(2), values[2]);
570 }
571
572 #[test]
573 fn test_try_new_from_builder() {
574 let values = vec![[1, 2, 3], [5, 6, 7], [6, 7, 8]];
575 _test_try_new_from_builder_generic_for_key_types::<UInt8Type, UInt16Type>(values.clone());
577 _test_try_new_from_builder_generic_for_key_types::<UInt16Type, UInt8Type>(values.clone());
579 _test_try_new_from_builder_generic_for_key_types::<Int8Type, Int16Type>(values.clone());
581 _test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type>(values.clone());
583 _test_try_new_from_builder_generic_for_key_types::<UInt8Type, Int16Type>(values.clone());
585 _test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt8Type>(values.clone());
586 _test_try_new_from_builder_generic_for_key_types::<Int8Type, UInt16Type>(values.clone());
587 _test_try_new_from_builder_generic_for_key_types::<Int32Type, Int16Type>(values.clone());
588 }
589
590 #[test]
591 fn test_try_new_from_builder_cast_fails() {
592 let mut source_builder = FixedSizeBinaryDictionaryBuilder::<UInt16Type>::new(2);
593 for i in 0u16..257u16 {
594 source_builder.append_value(vec![(i >> 8) as u8, i as u8]);
595 }
596
597 let result =
600 FixedSizeBinaryDictionaryBuilder::<UInt8Type>::try_new_from_builder(source_builder);
601 assert!(result.is_err());
602 if let Err(e) = result {
603 assert!(matches!(e, ArrowError::CastError(_)));
604 assert_eq!(
605 e.to_string(),
606 "Cast error: Can't cast dictionary keys from source type UInt16 to type UInt8"
607 );
608 }
609 }
610
611 #[test]
612 fn test_finish_preserve_values() {
613 let mut builder = FixedSizeBinaryDictionaryBuilder::<Int32Type>::new(3);
615 builder.append_value("aaa");
616 builder.append_value("bbb");
617 builder.append_value("ccc");
618 let dict = builder.finish_preserve_values();
619 assert_eq!(dict.keys().values(), &[0, 1, 2]);
620 let values = dict
621 .downcast_dict::<FixedSizeBinaryArray>()
622 .unwrap()
623 .into_iter()
624 .collect::<Vec<_>>();
625 assert_eq!(
626 values,
627 vec![
628 Some(b"aaa".as_slice()),
629 Some(b"bbb".as_slice()),
630 Some(b"ccc".as_slice())
631 ]
632 );
633
634 builder.append_value("ddd");
636 builder.append_value("eee");
637 let dict2 = builder.finish_preserve_values();
638
639 assert_eq!(dict2.keys().values(), &[3, 4]);
642 let values = dict2
643 .downcast_dict::<FixedSizeBinaryArray>()
644 .unwrap()
645 .into_iter()
646 .collect::<Vec<_>>();
647 assert_eq!(values, [Some(b"ddd".as_slice()), Some(b"eee".as_slice())]);
648
649 let all_values = dict2
651 .values()
652 .as_any()
653 .downcast_ref::<FixedSizeBinaryArray>()
654 .unwrap()
655 .into_iter()
656 .collect::<Vec<_>>();
657 assert_eq!(
658 all_values,
659 [
660 Some(b"aaa".as_slice()),
661 Some(b"bbb".as_slice()),
662 Some(b"ccc".as_slice()),
663 Some(b"ddd".as_slice()),
664 Some(b"eee".as_slice())
665 ]
666 );
667 }
668}