1use crate::ZWeight;
2use crate::storage::file::{FilterStats, TouchedWindowCount};
3use crate::trace::BatchLocation;
4use crate::trace::cursor::Position;
5use crate::trace::ord::merge_batcher::MergeBatcher;
6use crate::{
7 DBData, DBWeight, NumEntries, Timestamp,
8 algebra::Lattice,
9 dynamic::{
10 DataTrait, DynDataTyped, DynPair, DynVec, DynWeightedPairs, Erase, Factory, LeanVec,
11 WeightTrait, WithFactory,
12 },
13 trace::{
14 Batch, BatchFactories, BatchReader, BatchReaderFactories, Builder, Cursor, DbspSerializer,
15 Deserializer,
16 layers::{
17 Cursor as TrieCursor, Layer, LayerCursor, LayerFactories, Leaf, LeafFactories,
18 OrdOffset, Trie,
19 },
20 },
21 utils::{ConsolidatePairedSlices, Tup2},
22};
23use feldera_storage::FileReader;
24use rand::Rng;
25use rkyv::{Archive, Deserialize, Serialize};
26use size_of::SizeOf;
27use std::any::TypeId;
28use std::fmt::{self, Debug, Display, Formatter};
29use std::sync::Arc;
30
31pub type VecValBatchLayer<K, V, T, R, O> = Layer<K, Layer<V, Leaf<DynDataTyped<T>, R>, O>, O>;
32
33pub struct VecValBatchFactories<K, V, T, R>
34where
35 K: DataTrait + ?Sized,
36 V: DataTrait + ?Sized,
37 R: WeightTrait + ?Sized,
38 T: Timestamp,
39{
40 layer_factories: LayerFactories<K, LayerFactories<V, LeafFactories<DynDataTyped<T>, R>>>,
41 item_factory: &'static dyn Factory<DynPair<K, V>>,
42 consolidate_weights: &'static dyn ConsolidatePairedSlices<DynDataTyped<T>, R>,
43 weighted_item_factory: &'static dyn Factory<DynPair<DynPair<K, V>, R>>,
44 weighted_items_factory: &'static dyn Factory<DynWeightedPairs<DynPair<K, V>, R>>,
45 weighted_vals_factory: &'static dyn Factory<DynWeightedPairs<V, R>>,
46 time_diffs_factory: &'static dyn Factory<DynWeightedPairs<DynDataTyped<T>, R>>,
47}
48
49impl<K, V, T, R> Clone for VecValBatchFactories<K, V, T, R>
50where
51 K: DataTrait + ?Sized,
52 V: DataTrait + ?Sized,
53 R: WeightTrait + ?Sized,
54 T: Timestamp,
55{
56 fn clone(&self) -> Self {
57 Self {
58 layer_factories: self.layer_factories.clone(),
59 item_factory: self.item_factory,
60 consolidate_weights: self.consolidate_weights,
61 weighted_item_factory: self.weighted_item_factory,
62 weighted_items_factory: self.weighted_items_factory,
63 weighted_vals_factory: self.weighted_vals_factory,
64 time_diffs_factory: self.time_diffs_factory,
65 }
66 }
67}
68
69unsafe impl<K, V, T, R> Send for VecValBatchFactories<K, V, T, R>
70where
71 K: DataTrait + ?Sized,
72 V: DataTrait + ?Sized,
73 R: WeightTrait + ?Sized,
74 T: Timestamp,
75{
76}
77
78impl<K, V, T, R> BatchReaderFactories<K, V, T, R> for VecValBatchFactories<K, V, T, R>
79where
80 K: DataTrait + ?Sized,
81 V: DataTrait + ?Sized,
82 R: WeightTrait + ?Sized,
83 T: Timestamp,
84{
85 fn new<KType, VType, RType>() -> Self
86 where
87 KType: DBData + Erase<K>,
88 VType: DBData + Erase<V>,
89 RType: DBWeight + Erase<R>,
90 {
91 Self {
92 layer_factories: LayerFactories::new::<KType>(LayerFactories::new::<VType>(
93 LeafFactories::new::<T, RType>(),
94 )),
95 item_factory: WithFactory::<Tup2<KType, VType>>::FACTORY,
96 consolidate_weights: <dyn ConsolidatePairedSlices<_, _>>::factory::<T, RType>(),
97 weighted_item_factory: WithFactory::<Tup2<Tup2<KType, VType>, RType>>::FACTORY,
98 weighted_items_factory:
99 WithFactory::<LeanVec<Tup2<Tup2<KType, VType>, RType>>>::FACTORY,
100 weighted_vals_factory: WithFactory::<LeanVec<Tup2<VType, RType>>>::FACTORY,
101 time_diffs_factory: WithFactory::<LeanVec<Tup2<T, RType>>>::FACTORY,
102 }
103 }
104
105 fn key_factory(&self) -> &'static dyn Factory<K> {
106 self.layer_factories.key
107 }
108
109 fn keys_factory(&self) -> &'static dyn Factory<DynVec<K>> {
110 self.layer_factories.keys
111 }
112
113 fn val_factory(&self) -> &'static dyn Factory<V> {
114 self.layer_factories.child.key
115 }
116
117 fn weight_factory(&self) -> &'static dyn Factory<R> {
118 self.layer_factories.child.child.diff
119 }
120}
121
122impl<K, V, T, R> BatchFactories<K, V, T, R> for VecValBatchFactories<K, V, T, R>
123where
124 K: DataTrait + ?Sized,
125 V: DataTrait + ?Sized,
126 R: WeightTrait + ?Sized,
127 T: Timestamp,
128{
129 fn item_factory(&self) -> &'static dyn Factory<DynPair<K, V>> {
132 self.item_factory
133 }
134
135 fn weighted_item_factory(&self) -> &'static dyn Factory<DynPair<DynPair<K, V>, R>> {
136 self.weighted_item_factory
137 }
138
139 fn weighted_items_factory(&self) -> &'static dyn Factory<DynWeightedPairs<DynPair<K, V>, R>> {
140 self.weighted_items_factory
141 }
142
143 fn weighted_vals_factory(&self) -> &'static dyn Factory<DynWeightedPairs<V, R>> {
144 self.weighted_vals_factory
145 }
146
147 fn time_diffs_factory(
148 &self,
149 ) -> Option<&'static dyn Factory<DynWeightedPairs<DynDataTyped<T>, R>>> {
150 Some(self.time_diffs_factory)
151 }
152
153 }
183
184pub struct VecValBatch<K, V, T, R, O = usize>
187where
188 K: DataTrait + ?Sized,
189 V: DataTrait + ?Sized,
190 R: WeightTrait + ?Sized,
191 T: Timestamp,
192 O: OrdOffset,
193{
194 factories: VecValBatchFactories<K, V, T, R>,
195
196 pub layer: VecValBatchLayer<K, V, T, R, O>,
204 touched_window_count: TouchedWindowCount,
205}
206
207impl<K, V, T, R, O> SizeOf for VecValBatch<K, V, T, R, O>
208where
209 K: DataTrait + ?Sized,
210 V: DataTrait + ?Sized,
211 R: WeightTrait + ?Sized,
212 T: Timestamp,
213 O: OrdOffset,
214{
215 fn size_of_children(&self, context: &mut size_of::Context) {
216 context.add(self.approximate_byte_size());
219 }
220}
221
222unsafe impl<K, V, T, R, O> Send for VecValBatch<K, V, T, R, O>
223where
224 K: DataTrait + ?Sized,
225 V: DataTrait + ?Sized,
226 R: WeightTrait + ?Sized,
227 T: Timestamp,
228 O: OrdOffset,
229{
230}
231
232impl<K, V, T, R, O> Debug for VecValBatch<K, V, T, R, O>
233where
234 K: DataTrait + ?Sized,
235 V: DataTrait + ?Sized,
236 R: WeightTrait + ?Sized,
237 T: Timestamp,
238 O: OrdOffset,
239{
240 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
241 f.debug_struct("VecValBatch")
242 .field("layer", &self.layer)
243 .finish()
244 }
245}
246
247impl<K, V, T: Lattice, R, O: OrdOffset> Deserialize<VecValBatch<K, V, T, R, O>, Deserializer> for ()
248where
249 K: DataTrait + ?Sized,
250 V: DataTrait + ?Sized,
251 R: WeightTrait + ?Sized,
252 T: Timestamp,
253 O: OrdOffset,
254{
255 fn deserialize(
256 &self,
257 _deserializer: &mut Deserializer,
258 ) -> Result<VecValBatch<K, V, T, R, O>, <Deserializer as rkyv::Fallible>::Error> {
259 todo!()
260 }
261}
262
263impl<K, V, T: Lattice, R, O: OrdOffset> Archive for VecValBatch<K, V, T, R, O>
264where
265 K: DataTrait + ?Sized,
266 V: DataTrait + ?Sized,
267 R: WeightTrait + ?Sized,
268 T: Timestamp,
269 O: OrdOffset,
270{
271 type Archived = ();
272 type Resolver = ();
273
274 unsafe fn resolve(&self, _pos: usize, _resolver: Self::Resolver, _out: *mut Self::Archived) {
275 todo!()
276 }
277}
278impl<K, V, T: Lattice, R, O: OrdOffset> Serialize<DbspSerializer<'_>> for VecValBatch<K, V, T, R, O>
279where
280 K: DataTrait + ?Sized,
281 V: DataTrait + ?Sized,
282 R: WeightTrait + ?Sized,
283 T: Timestamp,
284 O: OrdOffset,
285{
286 fn serialize(
287 &self,
288 _serializer: &mut DbspSerializer,
289 ) -> Result<Self::Resolver, <DbspSerializer<'_> as rkyv::Fallible>::Error> {
290 todo!()
291 }
292}
293
294impl<K, V, T, R, O> Clone for VecValBatch<K, V, T, R, O>
295where
296 K: DataTrait + ?Sized,
297 V: DataTrait + ?Sized,
298 R: WeightTrait + ?Sized,
299 T: Timestamp,
300 O: OrdOffset,
301{
302 fn clone(&self) -> Self {
303 Self {
304 factories: self.factories.clone(),
305 layer: self.layer.clone(),
306 touched_window_count: self.touched_window_count,
307 }
308 }
309}
310
311impl<K, V, T, R, O> NumEntries for VecValBatch<K, V, T, R, O>
312where
313 K: DataTrait + ?Sized,
314 V: DataTrait + ?Sized,
315 R: WeightTrait + ?Sized,
316 T: Timestamp,
317 O: OrdOffset,
318{
319 const CONST_NUM_ENTRIES: Option<usize> = <VecValBatchLayer<K, V, T, R, O>>::CONST_NUM_ENTRIES;
320
321 #[inline]
322 fn num_entries_shallow(&self) -> usize {
323 self.layer.num_entries_shallow()
324 }
325
326 #[inline]
327 fn num_entries_deep(&self) -> usize {
328 self.layer.num_entries_deep()
329 }
330}
331
332impl<K, V, T, R, O> Display for VecValBatch<K, V, T, R, O>
333where
334 K: DataTrait + ?Sized,
335 V: DataTrait + ?Sized,
336 R: WeightTrait + ?Sized,
337 T: Timestamp,
338 O: OrdOffset,
339{
340 fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
341 writeln!(
342 f,
343 "layer:\n{}",
344 textwrap::indent(&self.layer.to_string(), " ")
345 )
346 }
347}
348
349impl<K, V, T, R, O> BatchReader for VecValBatch<K, V, T, R, O>
350where
351 K: DataTrait + ?Sized,
352 V: DataTrait + ?Sized,
353 R: WeightTrait + ?Sized,
354 T: Timestamp,
355 O: OrdOffset,
356{
357 type Key = K;
358 type Val = V;
359 type Time = T;
360 type R = R;
361 type Factories = VecValBatchFactories<K, V, T, R>;
362
363 type Cursor<'s>
364 = VecValCursor<'s, K, V, T, R, O>
365 where
366 O: 's;
367
368 fn factories(&self) -> Self::Factories {
371 self.factories.clone()
372 }
373
374 fn cursor(&self) -> Self::Cursor<'_> {
375 VecValCursor {
376 cursor: self.layer.cursor(),
377 }
378 }
379
380 fn key_count(&self) -> usize {
385 <VecValBatchLayer<K, V, T, R, O> as Trie>::keys(&self.layer)
386 }
387
388 fn len(&self) -> usize {
389 <VecValBatchLayer<K, V, T, R, O> as Trie>::tuples(&self.layer)
390 }
391
392 fn approximate_byte_size(&self) -> usize {
393 self.layer.approximate_byte_size()
394 }
395
396 fn membership_filter_stats(&self) -> FilterStats {
397 FilterStats::default()
398 }
399
400 fn sample_keys<RG>(&self, rng: &mut RG, sample_size: usize, sample: &mut DynVec<Self::Key>)
401 where
402 RG: Rng,
403 {
404 self.layer.sample_keys(rng, sample_size, sample);
405 }
406}
407
408impl<K, V, T, R, O> Batch for VecValBatch<K, V, T, R, O>
409where
410 K: DataTrait + ?Sized,
411 V: DataTrait + ?Sized,
412 R: WeightTrait + ?Sized,
413 T: Timestamp,
414 O: OrdOffset,
415{
416 type Timed<T2: Timestamp> = VecValBatch<K, V, T2, R, O>;
417 type Batcher = MergeBatcher<Self>;
418 type Builder = VecValBuilder<K, V, T, R, O>;
419
420 fn file_reader(&self) -> Option<Arc<dyn FileReader>> {
421 unimplemented!()
422 }
423
424 fn key_bounds(&self) -> Option<(&Self::Key, &Self::Key)> {
425 Some((self.layer.keys.first()?, self.layer.keys.last()?))
426 }
427
428 fn negative_weight_count(&self) -> Option<u64> {
429 None
430 }
431
432 fn touched_window_count(&self) -> TouchedWindowCount {
433 self.touched_window_count
434 }
435}
436
437#[derive(Debug, SizeOf)]
439pub struct VecValCursor<'s, K, V, T, R, O = usize>
440where
441 K: DataTrait + ?Sized,
442 V: DataTrait + ?Sized,
443 R: WeightTrait + ?Sized,
444 T: Timestamp,
445 O: OrdOffset,
446{
447 cursor: LayerCursor<'s, K, Layer<V, Leaf<DynDataTyped<T>, R>, O>, O>,
448}
449
450impl<K, V, T, R, O> Clone for VecValCursor<'_, K, V, T, R, O>
451where
452 K: DataTrait + ?Sized,
453 V: DataTrait + ?Sized,
454 R: WeightTrait + ?Sized,
455 T: Timestamp,
456 O: OrdOffset,
457{
458 fn clone(&self) -> Self {
459 Self {
460 cursor: self.cursor.clone(),
461 }
462 }
463}
464
465impl<K, V, T, R, O> Cursor<K, V, T, R> for VecValCursor<'_, K, V, T, R, O>
466where
467 K: DataTrait + ?Sized,
468 V: DataTrait + ?Sized,
469 R: WeightTrait + ?Sized,
470 T: Timestamp,
471 O: OrdOffset,
472{
473 fn weight_factory(&self) -> &'static dyn Factory<R> {
474 self.cursor.child.child.storage.factories.diff
475 }
476
477 fn key(&self) -> &K {
478 self.cursor.item()
479 }
480
481 fn val(&self) -> &V {
482 self.cursor.child.item()
483 }
484
485 fn map_times(&mut self, logic: &mut dyn FnMut(&T, &R)) {
486 self.cursor.child.child.rewind();
487 while self.cursor.child.child.valid() {
488 logic(
489 self.cursor.child.child.current_key(),
490 self.cursor.child.child.current_diff(),
491 );
492 self.cursor.child.child.step();
493 }
494 }
495
496 fn map_times_through(&mut self, upper: &T, logic: &mut dyn FnMut(&T, &R)) {
497 self.cursor.child.child.rewind();
498 while self.cursor.child.child.valid() {
499 if self.cursor.child.child.item().0.less_equal(upper) {
500 logic(
501 self.cursor.child.child.current_key(),
502 self.cursor.child.child.current_diff(),
503 );
504 }
505 self.cursor.child.child.step();
506 }
507 }
508
509 fn map_values(&mut self, logic: &mut dyn FnMut(&V, &R))
510 where
511 T: PartialEq<()>,
512 {
513 while self.val_valid() {
514 logic(self.val(), self.cursor.child.child.current_diff());
515 self.step_val();
516 }
517 }
518
519 fn weight(&mut self) -> &R
520 where
521 T: PartialEq<()>,
522 {
523 self.weight_checked()
524 }
525
526 fn weight_checked(&mut self) -> &R {
527 if TypeId::of::<T>() == TypeId::of::<()>() {
528 debug_assert!(&self.cursor.child.child.valid());
529 self.cursor.child.child.current_diff()
530 } else {
531 panic!("VecValCursor::weight_checked called on non-unit timestamp type");
532 }
533 }
534
535 fn key_valid(&self) -> bool {
536 self.cursor.valid()
537 }
538 fn val_valid(&self) -> bool {
539 self.cursor.child.valid()
540 }
541 fn step_key(&mut self) {
542 self.cursor.step();
543 }
544 fn step_key_reverse(&mut self) {
545 self.cursor.step_reverse();
546 }
547 fn seek_key(&mut self, key: &K) {
548 self.cursor.seek(key);
549 }
550
551 fn seek_key_exact(&mut self, key: &K, _hash: Option<u64>) -> bool {
552 self.seek_key(key);
553 self.key_valid() && self.key().eq(key)
554 }
555
556 fn seek_key_with(&mut self, predicate: &dyn Fn(&K) -> bool) {
557 self.cursor.seek_with(predicate);
558 }
559 fn seek_key_with_reverse(&mut self, predicate: &dyn Fn(&K) -> bool) {
560 self.cursor.seek_with_reverse(predicate);
561 }
562 fn seek_key_reverse(&mut self, key: &K) {
563 self.cursor.seek_reverse(key);
564 }
565 fn step_val(&mut self) {
566 self.cursor.child.step();
567 }
568 fn seek_val(&mut self, val: &V) {
569 self.cursor.child.seek(val);
570 }
571 fn seek_val_with(&mut self, predicate: &dyn Fn(&V) -> bool) {
572 self.cursor.child.seek_with(predicate);
573 }
574 fn rewind_keys(&mut self) {
575 self.cursor.rewind();
576 }
577 fn fast_forward_keys(&mut self) {
578 self.cursor.fast_forward();
579 }
580 fn rewind_vals(&mut self) {
581 self.cursor.child.rewind();
582 }
583
584 fn step_val_reverse(&mut self) {
585 self.cursor.child.step_reverse();
586 }
587
588 fn seek_val_reverse(&mut self, val: &V) {
589 self.cursor.child.seek_reverse(val);
590 }
591
592 fn seek_val_with_reverse(&mut self, predicate: &dyn Fn(&V) -> bool) {
593 self.cursor.child.seek_with_reverse(predicate);
594 }
595
596 fn fast_forward_vals(&mut self) {
597 self.cursor.child.fast_forward();
598 }
599
600 fn position(&self) -> Option<Position> {
601 Some(Position {
602 total: TrieCursor::keys(&self.cursor) as u64,
603 offset: self.cursor.pos() as u64,
604 })
605 }
606}
607
608#[derive(SizeOf)]
610pub struct VecValBuilder<K, V, T, R, O = usize>
611where
612 K: DataTrait + ?Sized,
613 V: DataTrait + ?Sized,
614 R: WeightTrait + ?Sized,
615 T: Timestamp,
616 O: OrdOffset,
617{
618 #[size_of(skip)]
619 factories: VecValBatchFactories<K, V, T, R>,
620 keys: Box<DynVec<K>>,
621 offs: Vec<O>,
622 vals: Box<DynVec<V>>,
623 val_offs: Vec<O>,
624 times: Box<DynVec<DynDataTyped<T>>>,
625 diffs: Box<DynVec<R>>,
626 total_positive_weight: ZWeight,
627 total_negative_weight: ZWeight,
628}
629
630impl<K, V, T, R, O> VecValBuilder<K, V, T, R, O>
631where
632 K: DataTrait + ?Sized,
633 V: DataTrait + ?Sized,
634 R: WeightTrait + ?Sized,
635 T: Timestamp,
636 O: OrdOffset,
637{
638 fn pushed_key(&mut self) {
639 let off = O::from_usize(self.vals.len());
640 debug_assert!(off > *self.offs.last().unwrap());
641 self.offs.push(off);
642 }
643
644 fn pushed_val(&mut self) {
645 let val_off = O::from_usize(self.times.len());
646 debug_assert!(val_off > *self.val_offs.last().unwrap());
647 self.val_offs.push(val_off);
648 }
649}
650
651impl<K, V, T, R, O> Builder<VecValBatch<K, V, T, R, O>> for VecValBuilder<K, V, T, R, O>
652where
653 Self: SizeOf,
654 K: DataTrait + ?Sized,
655 V: DataTrait + ?Sized,
656 R: WeightTrait + ?Sized,
657 T: Timestamp,
658 O: OrdOffset,
659{
660 fn with_capacity_in_location(
661 factories: &VecValBatchFactories<K, V, T, R>,
662 key_capacity: usize,
663 value_capacity: usize,
664 _location: Option<BatchLocation>,
665 ) -> Self {
666 let mut keys = factories.layer_factories.keys.default_box();
667 keys.reserve_exact(key_capacity);
668
669 let mut offs = Vec::with_capacity(key_capacity + 1);
670 offs.push(O::zero());
671
672 let mut vals = factories.layer_factories.child.keys.default_box();
673 vals.reserve_exact(value_capacity);
674
675 let mut val_offs = Vec::with_capacity(value_capacity + 1);
676 val_offs.push(O::zero());
677
678 let mut times = factories.layer_factories.child.child.keys.default_box();
679 times.reserve_exact(value_capacity);
680
681 let mut diffs = factories.layer_factories.child.child.diffs.default_box();
682 diffs.reserve_exact(value_capacity);
683 Self {
684 factories: factories.clone(),
685 keys,
686 offs,
687 vals,
688 val_offs,
689 times,
690 diffs,
691 total_positive_weight: 0,
692 total_negative_weight: 0,
693 }
694 }
695
696 fn reserve(&mut self, additional: usize) {
697 self.keys.reserve(additional);
698 self.offs.reserve(additional);
699 self.vals.reserve(additional);
700 self.times.reserve(additional);
701 self.diffs.reserve(additional);
702 }
703
704 fn push_key(&mut self, key: &K) {
705 self.keys.push_ref(key);
706 self.pushed_key();
707 }
708
709 fn push_key_mut(&mut self, key: &mut K) {
710 self.keys.push_val(key);
711 self.pushed_key();
712 }
713
714 fn push_val(&mut self, val: &V) {
715 self.vals.push_ref(val);
716 self.pushed_val();
717 }
718
719 fn push_val_mut(&mut self, val: &mut V) {
720 self.vals.push_val(val);
721 self.pushed_val();
722 }
723
724 fn push_time_diff(&mut self, time: &T, weight: &R) {
725 debug_assert!(!weight.is_zero());
726 self.times.push(time.clone());
727 self.diffs.push_ref(weight);
728 }
729
730 fn push_time_diff_mut(&mut self, time: &mut T, weight: &mut R) {
731 debug_assert!(!weight.is_zero());
732 self.times.push(time.clone());
733 self.diffs.push_val(weight);
734 }
735
736 fn done(self) -> VecValBatch<K, V, T, R, O> {
737 VecValBatch {
738 layer: Layer::from_parts(
739 &self.factories.layer_factories,
740 self.keys,
741 self.offs,
742 Layer::from_parts(
743 &self.factories.layer_factories.child,
744 self.vals,
745 self.val_offs,
746 Leaf::from_parts(
747 &self.factories.layer_factories.child.child,
748 self.times,
749 self.diffs,
750 ),
751 ),
752 ),
753 factories: self.factories,
754 touched_window_count: TouchedWindowCount::default(),
755 }
756 }
757
758 fn num_keys(&self) -> usize {
759 self.keys.len()
760 }
761
762 fn num_tuples(&self) -> usize {
763 self.diffs.len()
764 }
765}
766
767