dbsp 0.329.0

use crate::{
    Circuit, DBData, Stream, Timestamp, ZWeight,
    algebra::{AddAssignByRef, HasOne, HasZero, IndexedZSet, PartialOrder, ZTrace},
    circuit::{
        OwnershipPreference, Scope, WithClock,
        circuit_builder::register_replay_stream,
        metadata::{BatchSizeStats, INPUT_BATCHES_STATS, OUTPUT_BATCHES_STATS, OperatorMeta},
        operator_traits::{BinaryOperator, Operator},
    },
    dynamic::{ClonableTrait, DataTrait, DynOpt, DynPairs, Erase},
    operator::dynamic::{
        accumulate_trace::{
            AccumulateBoundsId, AccumulateTraceAppend, AccumulateTraceId, AccumulateZ1Trace,
            TimedSpine,
        },
        trace::{DelayedTraceId, TraceBounds},
    },
    trace::{
        Batch, BatchFactories, BatchReader, BatchReaderFactories, Builder, Cursor, TupleBuilder,
    },
};
use std::{borrow::Cow, marker::PhantomData, ops::Neg};

pub struct UpsertFactories<T: Timestamp, B: IndexedZSet> {
    pub batch_factories: B::Factories,
    pub trace_factories: <T::TimedBatch<B> as BatchReader>::Factories,
}

impl<T: Timestamp, B: IndexedZSet> Clone for UpsertFactories<T, B> {
    fn clone(&self) -> Self {
        Self {
            batch_factories: self.batch_factories.clone(),
            trace_factories: self.trace_factories.clone(),
        }
    }
}

impl<T, B> UpsertFactories<T, B>
where
    T: Timestamp,
    B: Batch + IndexedZSet,
{
    pub fn new<KType, VType>() -> Self
    where
        KType: DBData + Erase<B::Key>,
        VType: DBData + Erase<B::Val>,
    {
        Self {
            batch_factories: BatchReaderFactories::new::<KType, VType, ZWeight>(),
            trace_factories: BatchReaderFactories::new::<KType, VType, ZWeight>(),
        }
    }
}

impl<C, K, V> Stream<C, Box<DynPairs<K, DynOpt<V>>>>
where
    K: DataTrait + ?Sized,
    V: DataTrait + ?Sized,
    C: Circuit,
{
    /// Convert a stream of upserts into a stream of updates.
    ///
    /// The input stream carries changes to a key/value map in the form of
    /// _upserts_.  An upsert assigns a new value to a key (or `None` to
    /// remove the key from the map) without explicitly removing the old
    /// value, if any.  Upserts are produced by some operators, including
    /// [`Stream::aggregate`].  The operator converts upserts
    /// into batches of updates, which is the input format of most DBSP
    /// operators.
    ///
    /// The operator assumes that the input vector is sorted by key and contains
    /// exactly one value per key.
    ///
    /// This is a stateful operator that internally maintains the trace of the
    /// collection.
    pub fn upsert<B>(
        &self,
        persistent_id: Option<&str>,
        factories: &UpsertFactories<<C as WithClock>::Time, B>,
    ) -> Stream<C, B>
    where
        B: IndexedZSet<Key = K, Val = V>,
    {
        let circuit = self.circuit();

        assert!(
            self.is_sharded(),
            "upsert operator applied to a non-sharded collection"
        );

        // We build the following circuit to implement the upsert semantics.
        // The collection is accumulated into a trace using integrator
        // (UntimedTraceAppend + Z1Trace = integrator).  The `Upsert` operator
        // evaluates each upsert command in the input stream against the trace
        // and computes a batch of updates to be added to the trace.
        //
        // ```text
        //                          ┌────────────────────────────►
        //                          │
        //                          │
        //  self        ┌──────┐    │        ┌───────────┐  trace
        // ────────────►│Upsert├────┴───────►│TraceAppend├────┐
        //              └──────┘   delta     └───────────┘    │
        //                 ▲                  ▲               │
        //                 │                  │               │
        //                 │                  │   ┌───────┐   │
        //                 └──────────────────┴───┤Z1Trace│◄──┘
        //                    delayed_trace       └───────┘
        // ```
        circuit.region("upsert", || {
            let bounds = <TraceBounds<K, V>>::unbounded();

            let (delayed_trace, z1feedback) = circuit.add_feedback_persistent(
                persistent_id
                    .map(|name| format!("{name}.integral"))
                    .as_deref(),
                AccumulateZ1Trace::new(
                    &factories.trace_factories,
                    &factories.batch_factories,
                    false,
                    circuit.root_scope(),
                    bounds.clone(),
                ),
            );
            delayed_trace.mark_sharded();

            let delta = circuit
                .add_binary_operator(
                    <Upsert<TimedSpine<B, C>, B, _>>::new(
                        &factories.batch_factories,
                        bounds.clone(),
                        circuit.clone(),
                    ),
                    &delayed_trace,
                    self,
                )
                .mark_distinct();
            delta.mark_sharded();
            let replay_stream = z1feedback.operator_mut().prepare_replay_stream(&delta);

            let trace = circuit.add_binary_operator_with_preference(
                <AccumulateTraceAppend<TimedSpine<B, C>, B, C>>::new(
                    &factories.trace_factories,
                    circuit.clone(),
                ),
                (&delayed_trace, OwnershipPreference::STRONGLY_PREFER_OWNED),
                (
                    &delta
                        .dyn_accumulate(&factories.batch_factories)
                        .into_enabled_stream(),
                    OwnershipPreference::PREFER_OWNED,
                ),
            );
            trace.mark_sharded();

            z1feedback.connect_with_preference(&trace, OwnershipPreference::STRONGLY_PREFER_OWNED);

            register_replay_stream(circuit, &delta, &replay_stream, &factories.batch_factories);

            circuit.cache_insert(DelayedTraceId::new(trace.stream_id()), delayed_trace);
            circuit.cache_insert(AccumulateTraceId::new(delta.stream_id()), trace);
            circuit.cache_insert(AccumulateBoundsId::<B>::new(delta.stream_id()), bounds);
            delta
        })
    }
}

pub struct Upsert<T, B, C>
where
    B: Batch,
    T: BatchReader,
{
    batch_factories: B::Factories,
    clock: C,
    bounds: TraceBounds<T::Key, T::Val>,

    // Input batch sizes.
    input_batch_stats: BatchSizeStats,

    // Output batch sizes.
    output_batch_stats: BatchSizeStats,

    phantom: PhantomData<B>,
}

impl<T, B, C> Upsert<T, B, C>
where
    B: Batch,
    T: BatchReader,
{
    pub fn new(
        batch_factories: &B::Factories,
        bounds: TraceBounds<T::Key, T::Val>,
        clock: C,
    ) -> Self {
        Self {
            batch_factories: batch_factories.clone(),
            clock,
            bounds,
            input_batch_stats: BatchSizeStats::new(),
            output_batch_stats: BatchSizeStats::new(),
            phantom: PhantomData,
        }
    }
}

impl<T, B, C> Operator for Upsert<T, B, C>
where
    T: BatchReader,
    B: Batch,
    C: 'static,
{
    fn name(&self) -> Cow<'static, str> {
        Cow::from("Upsert")
    }

    fn metadata(&self, meta: &mut OperatorMeta) {
        meta.extend(metadata! {
            INPUT_BATCHES_STATS => self.input_batch_stats.metadata(),
            OUTPUT_BATCHES_STATS => self.output_batch_stats.metadata(),
        });
    }

    fn fixedpoint(&self, _scope: Scope) -> bool {
        true
    }
}

impl<T, B, C> BinaryOperator<T, Box<DynPairs<T::Key, DynOpt<T::Val>>>, B> for Upsert<T, B, C>
where
    T: ZTrace,
    B: IndexedZSet<Key = T::Key, Val = T::Val>,
    C: WithClock<Time = T::Time> + 'static,
{
    async fn eval(&mut self, trace: &T, updates: &Box<DynPairs<T::Key, DynOpt<T::Val>>>) -> B {
        // Inputs must be sorted by key
        debug_assert!(updates.is_sorted_by(&|u1, u2| u1.fst().cmp(u2.fst())));
        self.input_batch_stats.add_batch(updates.len());

        // ... and contain a single update per key.
        // TODO: implement this check.
        // debug_assert!(updates
        //     .windows(2)
        //     .all(|upserts| upserts[0].0 != upserts[1].0));

        let mut key_updates = self.batch_factories.weighted_items_factory().default_box();
        let mut item = self.batch_factories.weighted_item_factory().default_box();

        let mut trace_cursor = trace.cursor();

        let builder =
            B::Builder::with_capacity(&self.batch_factories, updates.len(), updates.len() * 2);
        let mut builder = TupleBuilder::new(&self.batch_factories, builder);

        // let val_filter = self.bounds.effective_val_filter();
        let key_filter = self.bounds.effective_key_filter();

        for kv in updates.dyn_iter() {
            let (key, val) = kv.split();

            if let Some(key_filter) = &key_filter
                && !(key_filter.filter_func())(key)
            {
                continue;
            }

            if let Some(val) = val.get() {
                // We used to have this check to filter out records before inserting them.
                // It doesn't work for LastN filters, plus I'm not sure it's useful anyway,
                // since GC normally happens in the background.
                // if let Some(val_filter) = &val_filter
                //     && !(val_filter.filter_func())(val)
                // {
                //     continue;
                // }

                let (kv, weight) = item.split_mut();
                let (k, v) = kv.split_mut();

                key.clone_to(k);
                val.clone_to(v);
                **weight = HasOne::one();

                key_updates.push_val(&mut *item);
            }

            if trace_cursor.seek_key_exact(key, None) {
                // println!("{}: found key in trace_cursor", Runtime::worker_index());
                while trace_cursor.val_valid() {
                    let mut weight = ZWeight::zero();
                    trace_cursor.map_times(&mut |t, w| {
                        if t.less_equal(&self.clock.time()) {
                            weight.add_assign_by_ref(w);
                        };
                    });

                    if !weight.is_zero() {
                        let (kv, w) = item.split_mut();
                        let (k, v) = kv.split_mut();

                        key.clone_to(k);
                        trace_cursor.val().clone_to(v);
                        **w = weight.neg();

                        key_updates.push_val(&mut *item);
                    }

                    trace_cursor.step_val();
                }
            }

            key_updates.consolidate();
            builder.extend(key_updates.dyn_iter_mut());
            key_updates.clear();
        }

        let result = builder.done();
        self.output_batch_stats.add_batch(result.len());
        result
    }

    fn input_preference(&self) -> (OwnershipPreference, OwnershipPreference) {
        (
            OwnershipPreference::PREFER_OWNED,
            OwnershipPreference::PREFER_OWNED,
        )
    }
}