diff --git a/crates/dbsp/src/algebra/zset.rs b/crates/dbsp/src/algebra/zset.rs index b4c1b11db22..21a93fee1d4 100644 --- a/crates/dbsp/src/algebra/zset.rs +++ b/crates/dbsp/src/algebra/zset.rs @@ -250,6 +250,11 @@ pub trait ZSet: IndexedZSet { /// negative, so the result can be zero even if the Z-set contains nonzero /// weights. fn weighted_count(&self, sum: &mut Self::R); + + /// Returns a Z-set that contains all elements with positive weights from + /// `self` with their weights preserved. + #[cfg(test)] + fn positive(&self) -> Self; } impl ZSet for Z @@ -266,6 +271,24 @@ where cursor.step_key(); } } + + #[cfg(test)] + fn positive(&self) -> Self { + let factories = self.factories(); + let mut builder = Self::Builder::with_capacity(&factories, self.key_count(), self.len()); + let mut cursor = self.cursor(); + + while cursor.key_valid() { + let weight = **cursor.weight(); + if weight > 0 { + builder.push_val_diff(cursor.val(), weight.erase()); + builder.push_key(cursor.key()); + } + cursor.step_key(); + } + + builder.done() + } } #[cfg(test)] diff --git a/crates/dbsp/src/circuit/circuit_builder.rs b/crates/dbsp/src/circuit/circuit_builder.rs index fdc1a7937a6..91f59bccf5a 100644 --- a/crates/dbsp/src/circuit/circuit_builder.rs +++ b/crates/dbsp/src/circuit/circuit_builder.rs @@ -544,6 +544,10 @@ dyn_clone::clone_trait_object!(StreamMetadata); /// data. It sets each record's weight to 1 if it is positive and drops the /// others. /// +/// The "positive" operator on a Z-set keeps positive weights unchanged and +/// maps all other weights to 0. It differs from "distinct" only in that it +/// does not clamp positive weights to 1. +/// /// ## Join on equal keys /// /// A DBSP equi-join takes batches `a` and `b` as input, finds all pairs of a diff --git a/crates/dbsp/src/mono.rs b/crates/dbsp/src/mono.rs index 28a7f416e7f..ae72ce47ef2 100644 --- a/crates/dbsp/src/mono.rs +++ b/crates/dbsp/src/mono.rs @@ -762,6 +762,13 @@ where self.inner().dyn_distinct_mono(&factories).typed() } + #[track_caller] + pub fn positive(&self) -> Self { + let factories = DistinctFactories::new::(); + + self.inner().dyn_positive_mono(&factories).typed() + } + #[track_caller] pub fn hash_distinct(&self) -> Self { let factories = HashDistinctFactories::new::(); @@ -1244,6 +1251,13 @@ where self.inner().dyn_distinct_mono(&factories).typed() } + #[track_caller] + pub fn positive(&self) -> Self { + let factories = DistinctFactories::new::(); + + self.inner().dyn_positive_mono(&factories).typed() + } + #[track_caller] pub fn hash_distinct(&self) -> Self { let factories = HashDistinctFactories::new::(); diff --git a/crates/dbsp/src/operator/distinct.rs b/crates/dbsp/src/operator/distinct.rs index 320429d9142..e5a7a1cab10 100644 --- a/crates/dbsp/src/operator/distinct.rs +++ b/crates/dbsp/src/operator/distinct.rs @@ -22,14 +22,7 @@ where self.inner().dyn_stream_distinct(&factories).typed() } -} -impl Stream -where - C: Circuit, - Z: IndexedZSet, - Z::InnerBatch: Send, -{ /// Incrementally deduplicate input stream. /// /// This is an incremental version of the @@ -62,4 +55,25 @@ where self.inner().dyn_has_distinct(&factories).typed() } + + /// Incrementally filter out tuples with non-positive weights. + /// + /// Given a stream of changes to relation `A`, computes a stream of + /// changes to relation `A'` that contains every `(key, weight)` tuple of + /// `A` with `weight > 0`, with the weight unchanged. + /// + /// This operator is like [`distinct`](`Self::distinct`), except that it + /// preserves the weights of the retained tuples instead of clamping them + /// to 1. + #[cfg(not(feature = "backend-mode"))] + #[track_caller] + pub fn positive(&self) -> Stream + where + Z: crate::typed_batch::ZSet, + { + let factories = + crate::operator::dynamic::distinct::DistinctFactories::new::(); + + self.inner().dyn_positive(&factories).typed() + } } diff --git a/crates/dbsp/src/operator/dynamic/distinct.rs b/crates/dbsp/src/operator/dynamic/distinct.rs index 3187b11ee89..0993728c83b 100644 --- a/crates/dbsp/src/operator/dynamic/distinct.rs +++ b/crates/dbsp/src/operator/dynamic/distinct.rs @@ -14,7 +14,7 @@ use crate::{ DBData, Runtime, Timestamp, ZWeight, algebra::{ AddByRef, HasOne, HasZero, IndexedZSet, Lattice, OrdIndexedZSet, OrdIndexedZSetFactories, - PartialOrder, ZRingValue, + PartialOrder, }, circuit::{ Circuit, Scope, Stream, WithClock, @@ -49,6 +49,45 @@ use super::{MonoIndexedZSet, MonoZSet}; circuit_cache_key!(DistinctId(StreamId => Stream)); circuit_cache_key!(DistinctIncrementalId(StreamId => Stream)); +circuit_cache_key!(PositiveIncrementalId(StreamId => Stream)); + +/// Abstract semantics of the generic `distinct` operator. +pub trait DistinctSemantics: 'static { + /// Name of the incremental operator specialized for the root scope. + const ROOT_SCOPE_NAME: &'static str; + /// Name of the incremental operator for nested scopes. + const NESTED_SCOPE_NAME: &'static str; + + /// Maps the weight `w` of a tuple in the input to the weight of + /// the same tuple in the output. + fn output_weight(w: ZWeight) -> ZWeight; +} + +/// [`DistinctSemantics`] of the `distinct` operator: tuples with positive +/// weight get weight 1. +pub struct ZeroOrOneWeight; + +impl DistinctSemantics for ZeroOrOneWeight { + const ROOT_SCOPE_NAME: &'static str = "DistinctIncrementalTotal"; + const NESTED_SCOPE_NAME: &'static str = "DistinctIncremental"; + + fn output_weight(w: ZWeight) -> ZWeight { + if w > 0 { 1 } else { 0 } + } +} + +/// [`DistinctSemantics`] of the `positive` operator: tuples with positive +/// weight keep their weight. +pub struct PositiveWeight; + +impl DistinctSemantics for PositiveWeight { + const ROOT_SCOPE_NAME: &'static str = "PositiveIncrementalTotal"; + const NESTED_SCOPE_NAME: &'static str = "PositiveIncremental"; + + fn output_weight(w: ZWeight) -> ZWeight { + w.max(0) + } +} pub struct DistinctFactories { pub input_factories: Z::Factories, @@ -189,6 +228,13 @@ impl Stream { self.dyn_distinct(factories) } + pub fn dyn_positive_mono( + &self, + factories: &DistinctFactories, + ) -> Stream { + self.dyn_positive(factories) + } + pub fn dyn_hash_distinct_mono( &self, factories: &HashDistinctFactories, @@ -221,6 +267,13 @@ impl Stream { self.dyn_distinct(factories) } + pub fn dyn_positive_mono( + &self, + factories: &DistinctFactories::Time>, + ) -> Stream { + self.dyn_positive(factories) + } + pub fn dyn_hash_distinct_mono( &self, factories: &HashDistinctFactories::Time>, @@ -307,9 +360,42 @@ where }) } + /// See [`Stream::positive`]. + pub fn dyn_positive(&self, factories: &DistinctFactories) -> Stream + where + Z: IndexedZSet + Send, + { + let circuit = self.circuit(); + circuit.region("positive", || { + let stream = self.try_sharded_version(); + + circuit + .cache_get_or_insert_with(PositiveIncrementalId::new(stream.stream_id()), || { + stream.dyn_positive_inner(factories) + }) + .clone() + }) + } + pub fn dyn_distinct_inner(&self, factories: &DistinctFactories) -> Stream where Z: IndexedZSet + Send, + { + self.distinct_inner_generic::(factories) + .mark_distinct() + } + + pub fn dyn_positive_inner(&self, factories: &DistinctFactories) -> Stream + where + Z: IndexedZSet + Send, + { + self.distinct_inner_generic::(factories) + } + + fn distinct_inner_generic(&self, factories: &DistinctFactories) -> Stream + where + Z: IndexedZSet + Send, + S: DistinctSemantics, { let circuit = self.circuit(); @@ -318,7 +404,7 @@ where if circuit.root_scope() == 0 { // Use an implementation optimized to work in the root scope. circuit.add_binary_operator( - StreamingBinaryWrapper::new(DistinctIncrementalTotal::new( + StreamingBinaryWrapper::new(DistinctIncrementalTotal::<_, _, S>::new( Location::caller(), &factories.input_factories, )), @@ -339,7 +425,7 @@ where // └───────────┘ └─────┘ └─────┘ └───────────────────┘ // ``` circuit.add_binary_operator( - StreamingBinaryWrapper::new(DistinctIncremental::new( + StreamingBinaryWrapper::new(DistinctIncremental::<_, _, _, S>::new( Location::caller(), &factories.input_factories, &factories.aux_factories, @@ -356,7 +442,6 @@ where ) } .mark_sharded() - .mark_distinct() } } @@ -403,14 +488,15 @@ where } } -/// Incremental version of the distinct operator that only works in the -/// top-level scope (i.e., for totally ordered timestamps). +/// Incremental version of the distinct and positive operators that only works +/// in the top-level scope (i.e., for totally ordered timestamps). /// /// Takes a stream `a` of changes to relation `A` and a stream with delayed /// value of `A`: `z^-1(A) = a.integrate().delay()` and computes -/// `distinct(A) - distinct(z^-1(A))` incrementally, by only considering -/// values in the support of `a`. -struct DistinctIncrementalTotal { +/// `f(A) - f(z^-1(A))` incrementally, by only considering values in the +/// support of `a`, where `f` applies [`DistinctSemantics::output_weight`] to +/// the weight of every tuple. +struct DistinctIncrementalTotal { input_factories: Z::Factories, location: &'static Location<'static>, @@ -423,10 +509,10 @@ struct DistinctIncrementalTotal { chunk_size: usize, first_chunk_size: usize, - _type: PhantomData<(Z, I)>, + _type: PhantomData<(Z, I, S)>, } -impl DistinctIncrementalTotal { +impl DistinctIncrementalTotal { pub fn new(location: &'static Location<'static>, input_factories: &Z::Factories) -> Self { Self { input_factories: input_factories.clone(), @@ -470,13 +556,14 @@ impl DistinctIncrementalTotal { } } -impl Operator for DistinctIncrementalTotal +impl Operator for DistinctIncrementalTotal where Z: IndexedZSet, I: 'static, + S: DistinctSemantics, { fn name(&self) -> Cow<'static, str> { - Cow::from("DistinctIncrementalTotal") + Cow::from(S::ROOT_SCOPE_NAME) } fn location(&self) -> OperatorLocation { @@ -495,10 +582,11 @@ where } } -impl StreamingBinaryOperator>, I, Z> for DistinctIncrementalTotal +impl StreamingBinaryOperator>, I, Z> for DistinctIncrementalTotal where Z: IndexedZSet, I: WithSnapshot + 'static, + S: DistinctSemantics, { fn eval( self: Rc, @@ -555,15 +643,13 @@ where let new_weight = old_weight.add_by_ref(&w); - if old_weight.le0() { - // Weight changes from non-positive to positive. - if new_weight.ge0() && !new_weight.is_zero() { - builder.push_val_diff(v, ZWeight::one().erase()); - any_values = true; - } - } else if new_weight.le0() { - // Weight changes from positive to non-positive. - builder.push_val_diff(v, ZWeight::one().neg().erase()); + // Emit the change in the tuple's output weight. For + // `distinct` it is nonzero only when the integrated + // weight changes sign; for `positive` also when a + // positive weight changes value. + let diff = S::output_weight(new_weight) - S::output_weight(old_weight); + if !diff.is_zero() { + builder.push_val_diff(v, diff.erase()); any_values = true; } @@ -574,12 +660,15 @@ where delta_cursor.step_val(); } } else { + // The key is missing from the integral: the weight + // if copied from the delta while delta_cursor.val_valid() { let new_weight = **delta_cursor.weight(); debug_assert!(!new_weight.is_zero()); - if new_weight.ge0() { - builder.push_val_diff(delta_cursor.val(), ZWeight::one().erase()); + let output_weight = S::output_weight(new_weight); + if !output_weight.is_zero() { + builder.push_val_diff(delta_cursor.val(), output_weight.erase()); any_values = true; } @@ -612,7 +701,7 @@ where type KeysOfInterest = BTreeMap, R>>>; #[derive(SizeOf)] -struct DistinctIncremental +struct DistinctIncremental where Z: IndexedZSet, T: WithSnapshot, @@ -649,15 +738,16 @@ where chunk_size: usize, first_chunk_size: usize, - _type: PhantomData<(Z, T)>, + _type: PhantomData<(Z, T, S)>, } -impl DistinctIncremental +impl DistinctIncremental where Z: IndexedZSet, T: WithSnapshot, T::Batch: ZBatchReader, Clk: WithClock