compute: allow specification of join yielding

This commit introduces the plumbing required to allow users of Compute
to specify the yielding behavior of linear join operators via
`UpdateConfiguration` commands.

Author: teskje

src/adapter/src/flags.rs

@@ -26,6 +26,7 @@ pub fn compute_config(config: &SystemVars) -> ComputeParameters {
     ComputeParameters {
         max_result_size: Some(config.max_result_size()),
         dataflow_max_inflight_bytes: Some(config.dataflow_max_inflight_bytes()),
+        linear_join_yielding: None,
         enable_mz_join_core: Some(config.enable_mz_join_core()),
         enable_jemalloc_profiling: Some(config.enable_jemalloc_profiling()),
         persist: persist_config(config),

src/compute-client/src/protocol/command.proto

@@ -69,4 +69,5 @@ message ProtoComputeParameters {
     mz_tracing.params.ProtoTracingParameters tracing = 5;
     mz_service.params.ProtoGrpcClientParameters grpc_client = 6;
     optional bool enable_jemalloc_profiling = 7;
+    mz_compute_types.dataflows.ProtoYieldSpec linear_join_yielding = 8;
 }

src/compute-client/src/protocol/command.rs

@@ -10,7 +10,7 @@
 //! Compute protocol commands.
 
 use mz_cluster_client::client::{ClusterStartupEpoch, TimelyConfig, TryIntoTimelyConfig};
-use mz_compute_types::dataflows::DataflowDescription;
+use mz_compute_types::dataflows::{DataflowDescription, YieldSpec};
 use mz_expr::RowSetFinishing;
 use mz_ore::tracing::OpenTelemetryContext;
 use mz_persist_client::cfg::PersistParameters;
@@ -359,6 +359,8 @@ pub struct ComputeParameters {
     pub max_result_size: Option<u32>,
     /// The maximum number of in-flight bytes emitted by persist_sources feeding dataflows.
     pub dataflow_max_inflight_bytes: Option<usize>,
+    /// The yielding behavior with which linear joins should be rendered.
+    pub linear_join_yielding: Option<YieldSpec>,
     /// Whether rendering should use `mz_join_core` rather than DD's `JoinCore::join_core`.
     pub enable_mz_join_core: Option<bool>,
     /// Whether to activate jemalloc heap profiling.
@@ -377,6 +379,7 @@ impl ComputeParameters {
         let ComputeParameters {
             max_result_size,
             dataflow_max_inflight_bytes,
+            linear_join_yielding,
             enable_mz_join_core,
             enable_jemalloc_profiling,
             persist,
@@ -390,6 +393,9 @@ impl ComputeParameters {
         if dataflow_max_inflight_bytes.is_some() {
             self.dataflow_max_inflight_bytes = dataflow_max_inflight_bytes;
         }
+        if linear_join_yielding.is_some() {
+            self.linear_join_yielding = linear_join_yielding;
+        }
         if enable_mz_join_core.is_some() {
             self.enable_mz_join_core = enable_mz_join_core;
         }
@@ -413,6 +419,7 @@ impl RustType<ProtoComputeParameters> for ComputeParameters {
         ProtoComputeParameters {
             max_result_size: self.max_result_size.into_proto(),
             dataflow_max_inflight_bytes: self.dataflow_max_inflight_bytes.into_proto(),
+            linear_join_yielding: self.linear_join_yielding.into_proto(),
             enable_mz_join_core: self.enable_mz_join_core.into_proto(),
             enable_jemalloc_profiling: self.enable_jemalloc_profiling.into_proto(),
             persist: Some(self.persist.into_proto()),
@@ -425,6 +432,7 @@ impl RustType<ProtoComputeParameters> for ComputeParameters {
         Ok(Self {
             max_result_size: proto.max_result_size.into_rust()?,
             dataflow_max_inflight_bytes: proto.dataflow_max_inflight_bytes.into_rust()?,
+            linear_join_yielding: proto.linear_join_yielding.into_rust()?,
             enable_mz_join_core: proto.enable_mz_join_core.into_rust()?,
             enable_jemalloc_profiling: proto.enable_jemalloc_profiling.into_rust()?,
             persist: proto

src/compute-types/src/dataflows.proto

@@ -13,6 +13,7 @@ import "compute-types/src/plan.proto";
 import "compute-types/src/sinks.proto";
 import "compute-types/src/sources.proto";
 import "expr/src/scalar.proto";
+import "proto/src/proto.proto";
 import "repr/src/antichain.proto";
 import "repr/src/global_id.proto";
 import "repr/src/relation_and_scalar.proto";
@@ -63,3 +64,10 @@ message ProtoBuildDesc {
     mz_repr.global_id.ProtoGlobalId id = 1;
     plan.ProtoPlan plan = 2;
 }
+
+message ProtoYieldSpec {
+    oneof kind {
+        uint64 by_work = 1;
+        mz_proto.ProtoDuration by_time = 2;
+    }
+}

src/compute-types/src/dataflows.rs

@@ -10,6 +10,7 @@
 //! Types for describing dataflows.
 
 use std::collections::{BTreeMap, BTreeSet};
+use std::time::Duration;
 
 use mz_expr::{CollectionPlan, MirRelationExpr, MirScalarExpr, OptimizedMirRelationExpr};
 use mz_proto::{IntoRustIfSome, ProtoMapEntry, ProtoType, RustType, TryFromProtoError};
@@ -640,6 +641,38 @@ impl RustType<ProtoBuildDesc> for BuildDesc<crate::plan::Plan> {
     }
 }
 
+/// Specification of a dataflow operator's yielding behavior.
+#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize, Arbitrary)]
+pub enum YieldSpec {
+    ByWork(usize),
+    ByTime(Duration),
+}
+
+impl RustType<ProtoYieldSpec> for YieldSpec {
+    fn into_proto(&self) -> ProtoYieldSpec {
+        use proto_yield_spec::Kind;
+
+        let kind = match *self {
+            Self::ByWork(w) => Kind::ByWork(w.into_proto()),
+            Self::ByTime(t) => Kind::ByTime(t.into_proto()),
+        };
+        ProtoYieldSpec { kind: Some(kind) }
+    }
+
+    fn from_proto(proto: ProtoYieldSpec) -> Result<Self, TryFromProtoError> {
+        use proto_yield_spec::Kind;
+
+        let Some(kind) = proto.kind else {
+            return Err(TryFromProtoError::missing_field("ProtoYieldSpec::kind"));
+        };
+        let spec = match kind {
+            Kind::ByWork(w) => Self::ByWork(w.into_rust()?),
+            Kind::ByTime(t) => Self::ByTime(t.into_rust()?),
+        };
+        Ok(spec)
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use mz_proto::protobuf_roundtrip;

src/compute/src/compute_state.rs

@@ -41,7 +41,7 @@ use crate::arrangement::manager::{TraceBundle, TraceManager};
 use crate::logging;
 use crate::logging::compute::ComputeEvent;
 use crate::metrics::ComputeMetrics;
-use crate::render::LinearJoinImpl;
+use crate::render::{LinearJoinImpl, LinearJoinSpec};
 use crate::server::ResponseSender;
 
 /// Worker-local state that is maintained across dataflows.
@@ -80,8 +80,8 @@ pub struct ComputeState {
     max_result_size: u32,
     /// Maximum number of in-flight bytes emitted by persist_sources feeding dataflows.
     pub dataflow_max_inflight_bytes: usize,
-    /// Implementation to use for rendering linear joins.
-    pub linear_join_impl: LinearJoinImpl,
+    /// Specification for rendering linear joins.
+    pub linear_join_spec: LinearJoinSpec,
     /// Metrics for this replica.
     pub metrics: ComputeMetrics,
     /// A process-global handle to tracing configuration.
@@ -110,7 +110,7 @@ impl ComputeState {
             command_history,
             max_result_size: u32::MAX,
             dataflow_max_inflight_bytes: usize::MAX,
-            linear_join_impl: Default::default(),
+            linear_join_spec: Default::default(),
             metrics,
             tracing_handle,
         }
@@ -191,6 +191,7 @@ impl<'a, A: Allocate + 'static> ActiveComputeState<'a, A> {
         let ComputeParameters {
             max_result_size,
             dataflow_max_inflight_bytes,
+            linear_join_yielding,
             enable_mz_join_core,
             enable_jemalloc_profiling,
             persist,
@@ -204,8 +205,11 @@ impl<'a, A: Allocate + 'static> ActiveComputeState<'a, A> {
         if let Some(v) = dataflow_max_inflight_bytes {
             self.compute_state.dataflow_max_inflight_bytes = v;
         }
+        if let Some(v) = linear_join_yielding {
+            self.compute_state.linear_join_spec.yielding = v;
+        }
         if let Some(v) = enable_mz_join_core {
-            self.compute_state.linear_join_impl = match v {
+            self.compute_state.linear_join_spec.implementation = match v {
                 false => LinearJoinImpl::DifferentialDataflow,
                 true => LinearJoinImpl::Materialize,
             };

src/compute/src/render/context.rs

@@ -38,7 +38,7 @@ use timely::progress::{Antichain, Timestamp};
 
 use crate::extensions::arrange::{KeyCollection, MzArrange};
 use crate::render::errors::ErrorLogger;
-use crate::render::join::LinearJoinImpl;
+use crate::render::join::LinearJoinSpec;
 use crate::typedefs::{ErrSpine, RowSpine, TraceErrHandle, TraceRowHandle};
 
 // Local type definition to avoid the horror in signatures.
@@ -94,8 +94,8 @@ where
     pub bindings: BTreeMap<Id, CollectionBundle<S, V, T>>,
     /// A token that operators can probe to know whether the dataflow is shutting down.
     pub(super) shutdown_token: ShutdownToken,
-    /// The implementation to use for rendering linear joins.
-    pub(super) linear_join_impl: LinearJoinImpl,
+    /// Specification for rendering linear joins.
+    pub(super) linear_join_spec: LinearJoinSpec,
 }
 
 impl<S: Scope, V: Data + columnation::Columnation> Context<S, V>
@@ -122,7 +122,7 @@ where
             until: dataflow.until.clone(),
             bindings: BTreeMap::new(),
             shutdown_token: Default::default(),
-            linear_join_impl: Default::default(),
+            linear_join_spec: Default::default(),
         }
     }
 }

src/compute/src/render/join/linear_join.rs

@@ -11,10 +11,13 @@
 //!
 //! Consult [LinearJoinPlan] documentation for details.
 
+use std::time::Instant;
+
 use differential_dataflow::lattice::Lattice;
 use differential_dataflow::operators::arrange::arrangement::Arranged;
 use differential_dataflow::trace::TraceReader;
 use differential_dataflow::{AsCollection, Collection, Data};
+use mz_compute_types::dataflows::YieldSpec;
 use mz_compute_types::plan::join::linear_join::{LinearJoinPlan, LinearStagePlan};
 use mz_compute_types::plan::join::JoinClosure;
 use mz_repr::{DatumVec, Diff, Row, RowArena};
@@ -34,16 +37,31 @@ use crate::typedefs::RowSpine;
 /// Available linear join implementations.
 ///
 /// See the `mz_join_core` module docs for our rationale for providing two join implementations.
-#[derive(Clone, Copy, Default)]
+#[derive(Clone, Copy)]
 pub enum LinearJoinImpl {
-    #[default]
     Materialize,
     DifferentialDataflow,
 }
 
-impl LinearJoinImpl {
-    /// Run this join implementation on the provided arrangements.
-    fn run<G, Tr1, Tr2, L, I>(
+/// Specification of how linear joins are to be executed.
+#[derive(Clone, Copy)]
+pub struct LinearJoinSpec {
+    pub implementation: LinearJoinImpl,
+    pub yielding: YieldSpec,
+}
+
+impl Default for LinearJoinSpec {
+    fn default() -> Self {
+        Self {
+            implementation: LinearJoinImpl::Materialize,
+            yielding: YieldSpec::ByWork(1_000_000),
+        }
+    }
+}
+
+impl LinearJoinSpec {
+    /// Render a join operator according to this specification.
+    fn render<G, Tr1, Tr2, L, I>(
         &self,
         arranged1: &Arranged<G, Tr1>,
         arranged2: &Arranged<G, Tr2>,
@@ -58,13 +76,21 @@ impl LinearJoinImpl {
         I: IntoIterator,
         I::Item: Data,
     {
-        match self {
-            Self::DifferentialDataflow => {
+        use LinearJoinImpl::*;
+        use YieldSpec::*;
+
+        match (self.implementation, self.yielding) {
+            (DifferentialDataflow, _) => {
                 differential_dataflow::operators::JoinCore::join_core(arranged1, arranged2, result)
             }
-            Self::Materialize => mz_join_core(arranged1, arranged2, result, |_start, work| {
-                work >= 1_000_000
-            }),
+            (Materialize, ByWork(limit)) => {
+                let yield_fn = move |_start, work| work >= limit;
+                mz_join_core(arranged1, arranged2, result, yield_fn)
+            }
+            (Materialize, ByTime(limit)) => {
+                let yield_fn = move |start: Instant, _work| start.elapsed() >= limit;
+                mz_join_core(arranged1, arranged2, result, yield_fn)
+            }
         }
     }
 }
@@ -157,7 +183,7 @@ where
                 // Different variants of `joined` implement this differently,
                 // and the logic is centralized there.
                 let stream = differential_join(
-                    self.linear_join_impl,
+                    self.linear_join_spec,
                     joined,
                     inputs[stage_plan.lookup_relation].enter_region(inner),
                     stage_plan,
@@ -208,7 +234,7 @@ where
 /// version of the join of previous inputs. This is split into its own method
 /// to enable reuse of code with different types of `prev_keyed`.
 fn differential_join<G, T>(
-    join_impl: LinearJoinImpl,
+    join_spec: LinearJoinSpec,
     mut joined: JoinedFlavor<G, T>,
     lookup_relation: CollectionBundle<G, Row, T>,
     LinearStagePlan {
@@ -263,27 +289,27 @@ where
         }
         JoinedFlavor::Local(local) => match arrangement {
             ArrangementFlavor::Local(oks, errs1) => {
-                let (oks, errs2) = differential_join_inner(join_impl, local, oks, closure);
+                let (oks, errs2) = differential_join_inner(join_spec, local, oks, closure);
                 errors.push(errs1.as_collection(|k, _v| k.clone()));
                 errors.extend(errs2);
                 oks
             }
             ArrangementFlavor::Trace(_gid, oks, errs1) => {
-                let (oks, errs2) = differential_join_inner(join_impl, local, oks, closure);
+                let (oks, errs2) = differential_join_inner(join_spec, local, oks, closure);
                 errors.push(errs1.as_collection(|k, _v| k.clone()));
                 errors.extend(errs2);
                 oks
             }
         },
         JoinedFlavor::Trace(trace) => match arrangement {
             ArrangementFlavor::Local(oks, errs1) => {
-                let (oks, errs2) = differential_join_inner(join_impl, trace, oks, closure);
+                let (oks, errs2) = differential_join_inner(join_spec, trace, oks, closure);
                 errors.push(errs1.as_collection(|k, _v| k.clone()));
                 errors.extend(errs2);
                 oks
             }
             ArrangementFlavor::Trace(_gid, oks, errs1) => {
-                let (oks, errs2) = differential_join_inner(join_impl, trace, oks, closure);
+                let (oks, errs2) = differential_join_inner(join_spec, trace, oks, closure);
                 errors.push(errs1.as_collection(|k, _v| k.clone()));
                 errors.extend(errs2);
                 oks
@@ -299,7 +325,7 @@ where
 /// The return type includes an optional error collection, which may be
 /// `None` if we can determine that `closure` cannot error.
 fn differential_join_inner<G, T, Tr1, Tr2>(
-    join_impl: LinearJoinImpl,
+    join_spec: LinearJoinSpec,
     prev_keyed: Arranged<G, Tr1>,
     next_input: Arranged<G, Tr2>,
     closure: JoinClosure,
@@ -319,8 +345,8 @@ where
     let mut row_builder = Row::default();
 
     if closure.could_error() {
-        let (oks, err) = join_impl
-            .run(&prev_keyed, &next_input, move |key, old, new| {
+        let (oks, err) = join_spec
+            .render(&prev_keyed, &next_input, move |key, old, new| {
                 let temp_storage = RowArena::new();
                 let mut datums_local = datums.borrow_with_many(&[key, old, new]);
                 closure
@@ -339,7 +365,7 @@ where
 
         (oks.as_collection(), Some(err.as_collection()))
     } else {
-        let oks = join_impl.run(&prev_keyed, &next_input, move |key, old, new| {
+        let oks = join_spec.render(&prev_keyed, &next_input, move |key, old, new| {
             let temp_storage = RowArena::new();
             let mut datums_local = datums.borrow_with_many(&[key, old, new]);
             closure

src/compute/src/render/join/mod.rs

@@ -15,4 +15,4 @@ mod delta_join;
 mod linear_join;
 mod mz_join_core;
 
-pub use linear_join::LinearJoinImpl;
+pub use linear_join::{LinearJoinImpl, LinearJoinSpec};