Fuel flat_map

misc/python/materialize/mzcompose/__init__.py

@@ -404,6 +404,7 @@ def get_default_system_parameters(
     "enable_columnar_lgalloc",
     "compute_server_maintenance_interval",
     "compute_dataflow_max_inflight_bytes_cc",
+    "compute_flat_map_fuel",
     "consolidating_vec_growth_dampener",
     "copy_to_s3_parquet_row_group_file_ratio",
     "copy_to_s3_arrow_builder_buffer_ratio",

misc/python/materialize/parallel_workload/action.py

@@ -1118,6 +1118,7 @@ def __init__(
             "compute_server_maintenance_interval",
             "compute_dataflow_max_inflight_bytes",
             "compute_dataflow_max_inflight_bytes_cc",
+            "compute_flat_map_fuel",
             "consolidating_vec_growth_dampener",
             "compute_hydration_concurrency",
             "copy_to_s3_parquet_row_group_file_ratio",

src/compute-types/src/dyncfgs.rs

@@ -218,6 +218,14 @@ pub const COMPUTE_APPLY_COLUMN_DEMANDS: Config<bool> = Config::new(
     "When enabled, passes applys column demands to the RelationDesc used to read out of Persist.",
 );
 
+/// The amount of output the flat-map operator produces before yielding. Set to a high value to
+/// avoid yielding, or to a low value to yield frequently.
+pub const COMPUTE_FLAT_MAP_FUEL: Config<usize> = Config::new(
+    "compute_flat_map_fuel",
+    1_000_000,
+    "The amount of output the flat-map operator produces before yielding.",
+);
+
 /// Whether to render `as_specific_collection` using a fueled flat-map operator.
 pub const ENABLE_COMPUTE_RENDER_FUELED_AS_SPECIFIC_COLLECTION: Config<bool> = Config::new(
     "enable_compute_render_fueled_as_specific_collection",
@@ -355,6 +363,7 @@ pub fn all_dyncfgs(configs: ConfigSet) -> ConfigSet {
         .add(&ENABLE_COMPUTE_REPLICA_EXPIRATION)
         .add(&COMPUTE_REPLICA_EXPIRATION_OFFSET)
         .add(&COMPUTE_APPLY_COLUMN_DEMANDS)
+        .add(&COMPUTE_FLAT_MAP_FUEL)
         .add(&CONSOLIDATING_VEC_GROWTH_DAMPENER)
         .add(&ENABLE_COMPUTE_RENDER_FUELED_AS_SPECIFIC_COLLECTION)
         .add(&ENABLE_COMPUTE_LOGICAL_BACKPRESSURE)

src/compute/src/render/flat_map.rs

@@ -8,6 +8,7 @@
 // by the Apache License, Version 2.0.
 
 use differential_dataflow::consolidation::ConsolidatingContainerBuilder;
+use mz_compute_types::dyncfgs::COMPUTE_FLAT_MAP_FUEL;
 use mz_expr::MfpPlan;
 use mz_expr::{MapFilterProject, MirScalarExpr, TableFunc};
 use mz_repr::{DatumVec, RowArena, SharedRow};
@@ -41,15 +42,26 @@ where
         let (ok_collection, err_collection) =
             input.as_specific_collection(input_key.as_deref(), &self.config_set);
         let stream = ok_collection.inner;
-        let (oks, errs) = stream.unary_fallible(Pipeline, "FlatMapStage", move |_, _| {
+        let scope = input.scope();
+
+        // Budget to limit the number of rows processed in a single invocation.
+        //
+        // The current implementation can only yield between input batches, but not from within
+        // a batch. A `generate_series` can still cause unavailability if it generates many rows.
+        let budget = COMPUTE_FLAT_MAP_FUEL.get(&self.config_set);
+
+        let (oks, errs) = stream.unary_fallible(Pipeline, "FlatMapStage", move |_, info| {
+            let activator = scope.activator_for(info.address);
             Box::new(move |input, ok_output, err_output| {
                 let mut datums = DatumVec::new();
                 let mut datums_mfp = DatumVec::new();
 
                 // Buffer for extensions to `input_row`.
                 let mut table_func_output = Vec::new();
 
-                input.for_each(|cap, data| {
+                let mut budget = budget;
+
+                while let Some((cap, data)) = input.next() {
                     let mut ok_session = ok_output.session_with_builder(&cap);
                     let mut err_session = err_output.session_with_builder(&cap);
 
@@ -92,11 +104,16 @@ where
                                 &until,
                                 &mut ok_session,
                                 &mut err_session,
+                                &mut budget,
                             );
                             table_func_output.clear();
                         }
                     }
-                })
+                    if budget == 0 {
+                        activator.activate();
+                        break;
+                    }
+                }
             })
         });
 
@@ -129,6 +146,7 @@ fn drain_through_mfp<T>(
         ConsolidatingContainerBuilder<Vec<(DataflowError, T, Diff)>>,
         Counter<T, Vec<(DataflowError, T, Diff)>, Tee<T, Vec<(DataflowError, T, Diff)>>>,
     >,
+    budget: &mut usize,
 ) where
     T: crate::render::RenderTimestamp,
 {
@@ -156,6 +174,7 @@ fn drain_through_mfp<T>(
         );
 
         for result in results {
+            *budget = budget.saturating_sub(1);
             match result {
                 Ok((row, event_time, diff)) => {
                     // Copy the whole time, and re-populate event time.