risingwave_frontend/optimizer/property/
distribution.rs

1// Copyright 2022 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//!   "A -> B" represent A satisfies B
16//!                                 x
17//!  only as a required property    x  can used as both required
18//!                                 x  and provided property
19//!                                 x
20//!            ┌───┐                x┌──────┐
21//!            │Any◄─────────────────┤single│
22//!            └─▲─┘                x└──────┘
23//!              │                  x
24//!              │                  x
25//!              │                  x
26//!          ┌───┴────┐             x┌──────────┐
27//!          │AnyShard◄──────────────┤SomeShard │
28//!          └───▲────┘             x└──────────┘
29//!              │                  x
30//!          ┌───┴───────────┐      x┌──────────────┐ ┌──────────────┐
31//!          │ShardByKey(a,b)◄───┬───┤HashShard(a,b)│ │HashShard(b,a)│
32//!          └───▲──▲────────┘   │  x└──────────────┘ └┬─────────────┘
33//!              │  │            │  x                  │
34//!              │  │            └─────────────────────┘
35//!              │  │               x
36//!              │ ┌┴────────────┐  x┌────────────┐
37//!              │ │ShardByKey(a)◄───┤HashShard(a)│
38//!              │ └─────────────┘  x└────────────┘
39//!              │                  x
40//!             ┌┴────────────┐     x┌────────────┐
41//!             │ShardByKey(b)◄──────┤HashShard(b)│
42//!             └─────────────┘     x└────────────┘
43//!                                 x
44//!                                 x
45use std::collections::HashMap;
46use std::fmt;
47use std::fmt::Debug;
48
49use fixedbitset::FixedBitSet;
50use generic::PhysicalPlanRef;
51use itertools::Itertools;
52use risingwave_batch::worker_manager::worker_node_manager::WorkerNodeSelector;
53use risingwave_common::catalog::{FieldDisplay, Schema, TableId};
54use risingwave_common::hash::WorkerSlotId;
55use risingwave_pb::batch_plan::ExchangeInfo;
56use risingwave_pb::batch_plan::exchange_info::{
57    ConsistentHashInfo, Distribution as PbDistribution, DistributionMode, HashInfo,
58};
59
60use super::super::plan_node::*;
61use crate::catalog::FragmentId;
62use crate::catalog::catalog_service::CatalogReader;
63use crate::error::Result;
64use crate::optimizer::property::Order;
65
66/// the distribution property provided by a operator.
67#[derive(Debug, Clone, PartialEq, Eq, Hash)]
68pub enum Distribution {
69    /// There is only one partition. All records are placed on it.
70    ///
71    /// Note: singleton will not be enforced automatically.
72    /// It's set in `crate::stream_fragmenter::build_fragment`,
73    /// by setting `requires_singleton` manually.
74    Single,
75    /// Records are sharded into partitions, and satisfy the `AnyShard` but without any guarantee
76    /// about their placement rules.
77    SomeShard,
78    /// Records are sharded into partitions based on the hash value of some columns, which means
79    /// the records with the same hash values must be on the same partition.
80    /// `usize` is the index of column used as the distribution key.
81    HashShard(Vec<usize>),
82    /// A special kind of provided distribution which is almost the same as
83    /// [`Distribution::HashShard`], but may have different vnode mapping.
84    ///
85    /// It exists because the upstream MV can be scaled independently. So we use
86    /// `UpstreamHashShard` to **force an exchange to be inserted**.
87    ///
88    /// Alternatively, [`Distribution::SomeShard`] can also be used to insert an exchange, but
89    /// `UpstreamHashShard` contains distribution keys, which might be useful in some cases, e.g.,
90    /// two-phase Agg. It also satisfies [`RequiredDist::ShardByKey`].
91    ///
92    /// `TableId` is used to represent the data distribution(`vnode_mapping`) of this
93    /// `UpstreamHashShard`. The scheduler can fetch `TableId`'s corresponding `vnode_mapping` to do
94    /// shuffle.
95    UpstreamHashShard(Vec<usize>, TableId),
96    /// Records are available on all downstream shards.
97    Broadcast,
98}
99
100/// the distribution property requirement.
101#[derive(Debug, Clone, PartialEq)]
102pub enum RequiredDist {
103    /// with any distribution
104    Any,
105    /// records are shard on partitions, which means every record should belong to a partition
106    AnyShard,
107    /// records are shard on partitions based on some keys(order-irrelevance, ShardByKey({a,b}) is
108    /// equivalent with ShardByKey({b,a})), which means the records with same keys must be on
109    /// the same partition, as required property only. Any distribution sharded by a subset of this
110    /// key set satisfies the requirement.
111    ShardByKey(FixedBitSet),
112    /// records are shard on partitions based on an exact set of keys (order-irrelevance).
113    /// Only distribution sharded by the same key set satisfies this requirement.
114    ShardByExactKey(FixedBitSet),
115    /// must be the same with the physical distribution
116    PhysicalDist(Distribution),
117}
118
119impl Distribution {
120    pub fn to_prost(
121        &self,
122        output_count: u32,
123        catalog_reader: &CatalogReader,
124        worker_node_manager: &WorkerNodeSelector,
125        batch_parallelism: usize,
126    ) -> Result<ExchangeInfo> {
127        let exchange_info = ExchangeInfo {
128            mode: match self {
129                Distribution::Single => DistributionMode::Single,
130                Distribution::HashShard(_) => DistributionMode::Hash,
131                // TODO: add round robin DistributionMode
132                Distribution::SomeShard => DistributionMode::Single,
133                Distribution::Broadcast => DistributionMode::Broadcast,
134                Distribution::UpstreamHashShard(_, _) => DistributionMode::ConsistentHash,
135            } as i32,
136            distribution: match self {
137                Distribution::Single => None,
138                Distribution::HashShard(key) => {
139                    assert!(
140                        !key.is_empty(),
141                        "hash key should not be empty, use `Single` instead"
142                    );
143                    Some(PbDistribution::HashInfo(HashInfo {
144                        output_count,
145                        key: key.iter().map(|num| *num as u32).collect(),
146                    }))
147                }
148                // TODO: add round robin distribution
149                Distribution::SomeShard => None,
150                Distribution::Broadcast => None,
151                Distribution::UpstreamHashShard(key, table_id) => {
152                    assert!(
153                        !key.is_empty(),
154                        "hash key should not be empty, use `Single` instead"
155                    );
156
157                    let vnode_mapping = worker_node_manager.fragment_mapping(
158                        Self::get_fragment_id(catalog_reader, *table_id)?,
159                        batch_parallelism,
160                    )?;
161
162                    let worker_slot_to_id_map: HashMap<WorkerSlotId, u32> = vnode_mapping
163                        .iter_unique()
164                        .enumerate()
165                        .map(|(i, worker_slot_id)| (worker_slot_id, i as u32))
166                        .collect();
167
168                    Some(PbDistribution::ConsistentHashInfo(ConsistentHashInfo {
169                        vmap: vnode_mapping
170                            .iter()
171                            .map(|id| worker_slot_to_id_map[&id])
172                            .collect_vec(),
173                        key: key.iter().map(|num| *num as u32).collect(),
174                    }))
175                }
176            },
177        };
178        Ok(exchange_info)
179    }
180
181    /// check if the distribution satisfies other required distribution
182    pub fn satisfies(&self, required: &RequiredDist) -> bool {
183        match required {
184            RequiredDist::Any => true,
185            RequiredDist::AnyShard => {
186                matches!(
187                    self,
188                    Distribution::SomeShard
189                        | Distribution::HashShard(_)
190                        | Distribution::UpstreamHashShard(_, _)
191                        | Distribution::Broadcast
192                )
193            }
194            RequiredDist::ShardByKey(required_key) => match self {
195                Distribution::HashShard(hash_key)
196                | Distribution::UpstreamHashShard(hash_key, _) => {
197                    hash_key.iter().all(|idx| required_key.contains(*idx))
198                }
199                _ => false,
200            },
201            RequiredDist::ShardByExactKey(required_key) => match self {
202                Distribution::HashShard(hash_key)
203                | Distribution::UpstreamHashShard(hash_key, _) => {
204                    hash_key.len() == required_key.count_ones(..)
205                        && hash_key.iter().all(|idx| required_key.contains(*idx))
206                }
207                _ => false,
208            },
209            RequiredDist::PhysicalDist(other) => self == other,
210        }
211    }
212
213    /// Get distribution column indices. Panics if the distribution is `SomeShard` or `Broadcast`.
214    pub fn dist_column_indices(&self) -> &[usize] {
215        self.dist_column_indices_opt()
216            .unwrap_or_else(|| panic!("cannot obtain distribution columns for {self:?}"))
217    }
218
219    /// Get distribution column indices. Returns `None` if the distribution is `SomeShard` or `Broadcast`.
220    pub fn dist_column_indices_opt(&self) -> Option<&[usize]> {
221        match self {
222            Distribution::Single => Some(&[]),
223            Distribution::HashShard(dists) | Distribution::UpstreamHashShard(dists, _) => {
224                Some(dists)
225            }
226            Distribution::SomeShard | Distribution::Broadcast => None,
227        }
228    }
229
230    #[inline(always)]
231    fn get_fragment_id(catalog_reader: &CatalogReader, table_id: TableId) -> Result<FragmentId> {
232        catalog_reader
233            .read_guard()
234            .get_any_table_by_id(table_id)
235            .map(|table| table.fragment_id)
236            .map_err(Into::into)
237    }
238}
239
240impl fmt::Display for Distribution {
241    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
242        f.write_str("[")?;
243        match self {
244            Self::Single => f.write_str("Single")?,
245            Self::SomeShard => f.write_str("SomeShard")?,
246            Self::Broadcast => f.write_str("Broadcast")?,
247            Self::HashShard(vec) | Self::UpstreamHashShard(vec, _) => {
248                for key in vec {
249                    std::fmt::Debug::fmt(&key, f)?;
250                }
251            }
252        }
253        f.write_str("]")
254    }
255}
256
257pub struct DistributionDisplay<'a> {
258    pub distribution: &'a Distribution,
259    pub input_schema: &'a Schema,
260}
261
262impl DistributionDisplay<'_> {
263    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
264        let that = self.distribution;
265        match that {
266            Distribution::Single => f.write_str("Single"),
267            Distribution::SomeShard => f.write_str("SomeShard"),
268            Distribution::Broadcast => f.write_str("Broadcast"),
269            Distribution::HashShard(vec) | Distribution::UpstreamHashShard(vec, _) => {
270                if let Distribution::HashShard(_) = that {
271                    f.write_str("HashShard(")?;
272                } else {
273                    f.write_str("UpstreamHashShard(")?;
274                }
275                for (pos, key) in vec.iter().copied().with_position() {
276                    std::fmt::Debug::fmt(
277                        &FieldDisplay(self.input_schema.fields.get(key).unwrap()),
278                        f,
279                    )?;
280                    match pos {
281                        itertools::Position::First | itertools::Position::Middle => {
282                            f.write_str(", ")?;
283                        }
284                        _ => {}
285                    }
286                }
287                f.write_str(")")
288            }
289        }
290    }
291}
292
293impl fmt::Debug for DistributionDisplay<'_> {
294    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
295        self.fmt(f)
296    }
297}
298
299impl fmt::Display for DistributionDisplay<'_> {
300    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
301        self.fmt(f)
302    }
303}
304
305impl RequiredDist {
306    pub fn single() -> Self {
307        Self::PhysicalDist(Distribution::Single)
308    }
309
310    pub fn shard_by_key(tot_col_num: usize, key: &[usize]) -> Self {
311        let mut cols = FixedBitSet::with_capacity(tot_col_num);
312        for i in key {
313            cols.insert(*i);
314        }
315        assert!(!cols.is_clear());
316        Self::ShardByKey(cols)
317    }
318
319    pub fn shard_by_exact_key(tot_col_num: usize, key: &[usize]) -> Self {
320        let mut cols = FixedBitSet::with_capacity(tot_col_num);
321        for i in key {
322            cols.insert(*i);
323        }
324        assert!(!cols.is_clear());
325        Self::ShardByExactKey(cols)
326    }
327
328    pub fn hash_shard(key: &[usize]) -> Self {
329        assert!(!key.is_empty());
330        Self::PhysicalDist(Distribution::HashShard(key.to_vec()))
331    }
332
333    pub fn batch_enforce_if_not_satisfies(
334        &self,
335        mut plan: BatchPlanRef,
336        required_order: &Order,
337    ) -> Result<BatchPlanRef> {
338        plan = required_order.enforce_if_not_satisfies(plan)?;
339        if !plan.distribution().satisfies(self) {
340            Ok(self.batch_enforce(plan, required_order))
341        } else {
342            Ok(plan)
343        }
344    }
345
346    pub fn streaming_enforce_if_not_satisfies(&self, plan: StreamPlanRef) -> Result<StreamPlanRef> {
347        if !plan.distribution().satisfies(self) {
348            Ok(self.stream_enforce(plan))
349        } else {
350            Ok(plan)
351        }
352    }
353
354    pub fn no_shuffle(plan: StreamPlanRef) -> StreamPlanRef {
355        StreamExchange::new_no_shuffle(plan).into()
356    }
357
358    /// check if the distribution satisfies other required distribution
359    pub fn satisfies(&self, required: &RequiredDist) -> bool {
360        match self {
361            RequiredDist::Any => matches!(required, RequiredDist::Any),
362            RequiredDist::AnyShard => {
363                matches!(required, RequiredDist::Any | RequiredDist::AnyShard)
364            }
365            RequiredDist::ShardByKey(key) => match required {
366                RequiredDist::Any | RequiredDist::AnyShard => true,
367                RequiredDist::ShardByKey(required_key) => key.is_subset(required_key),
368                RequiredDist::ShardByExactKey(required_key) => {
369                    key == required_key && key.count_ones(..) == 1
370                }
371                _ => false,
372            },
373            RequiredDist::ShardByExactKey(key) => match required {
374                RequiredDist::Any | RequiredDist::AnyShard => true,
375                RequiredDist::ShardByKey(required_key) => key.is_subset(required_key),
376                RequiredDist::ShardByExactKey(required_key) => key == required_key,
377                _ => false,
378            },
379            RequiredDist::PhysicalDist(dist) => dist.satisfies(required),
380        }
381    }
382
383    pub fn batch_enforce(&self, plan: BatchPlanRef, required_order: &Order) -> BatchPlanRef {
384        let dist = self.to_dist();
385        BatchExchange::new(plan, required_order.clone(), dist).into()
386    }
387
388    pub fn stream_enforce(&self, plan: StreamPlanRef) -> StreamPlanRef {
389        let dist = self.to_dist();
390        StreamExchange::new(plan, dist).into()
391    }
392
393    fn to_dist(&self) -> Distribution {
394        match self {
395            // all the distribution satisfy the Any, and the function can be only called by
396            // `enforce_if_not_satisfies`
397            RequiredDist::Any => unreachable!(),
398            // TODO: add round robin distributed type
399            RequiredDist::AnyShard => todo!(),
400            RequiredDist::ShardByKey(required_keys) => {
401                Distribution::HashShard(required_keys.ones().collect())
402            }
403            RequiredDist::ShardByExactKey(required_keys) => {
404                Distribution::HashShard(required_keys.ones().collect())
405            }
406            RequiredDist::PhysicalDist(dist) => dist.clone(),
407        }
408    }
409}
410
411impl StreamPlanRef {
412    /// Eliminate `SomeShard` distribution by using the stream key as the distribution key to
413    /// enforce the current plan to have a known distribution key.
414    pub fn enforce_concrete_distribution(self) -> Self {
415        match self.distribution() {
416            Distribution::SomeShard => {
417                RequiredDist::shard_by_key(self.schema().len(), self.expect_stream_key())
418                    .stream_enforce(self)
419            }
420            _ => self,
421        }
422    }
423}
424
425#[cfg(test)]
426mod tests {
427    use super::{Distribution, RequiredDist};
428
429    #[test]
430    fn hash_shard_satisfy() {
431        let d1 = Distribution::HashShard(vec![0, 1]);
432        let d2 = Distribution::HashShard(vec![1, 0]);
433        let d3 = Distribution::HashShard(vec![0]);
434        let d4 = Distribution::HashShard(vec![1]);
435
436        let r1 = RequiredDist::shard_by_key(2, &[0, 1]);
437        let r3 = RequiredDist::shard_by_key(2, &[0]);
438        let r4 = RequiredDist::shard_by_key(2, &[1]);
439        let r_exact = RequiredDist::shard_by_exact_key(2, &[0, 1]);
440        let r_exact_single = RequiredDist::shard_by_exact_key(2, &[0]);
441        assert!(d1.satisfies(&RequiredDist::PhysicalDist(d1.clone())));
442        assert!(d2.satisfies(&RequiredDist::PhysicalDist(d2.clone())));
443        assert!(d3.satisfies(&RequiredDist::PhysicalDist(d3.clone())));
444        assert!(d4.satisfies(&RequiredDist::PhysicalDist(d4.clone())));
445
446        assert!(!d2.satisfies(&RequiredDist::PhysicalDist(d1.clone())));
447        assert!(!d3.satisfies(&RequiredDist::PhysicalDist(d1.clone())));
448        assert!(!d4.satisfies(&RequiredDist::PhysicalDist(d1.clone())));
449
450        assert!(!d1.satisfies(&RequiredDist::PhysicalDist(d3.clone())));
451        assert!(!d2.satisfies(&RequiredDist::PhysicalDist(d3.clone())));
452        assert!(!d1.satisfies(&RequiredDist::PhysicalDist(d4.clone())));
453        assert!(!d2.satisfies(&RequiredDist::PhysicalDist(d4.clone())));
454
455        assert!(d1.satisfies(&r1));
456        assert!(d2.satisfies(&r1));
457        assert!(d3.satisfies(&r1));
458        assert!(d4.satisfies(&r1));
459
460        assert!(!d1.satisfies(&r3));
461        assert!(!d2.satisfies(&r3));
462        assert!(d3.satisfies(&r3));
463        assert!(!d4.satisfies(&r3));
464
465        assert!(!d1.satisfies(&r4));
466        assert!(!d2.satisfies(&r4));
467        assert!(!d3.satisfies(&r4));
468        assert!(d4.satisfies(&r4));
469
470        assert!(d1.satisfies(&r_exact));
471        assert!(d2.satisfies(&r_exact));
472        assert!(!d3.satisfies(&r_exact));
473        assert!(!d4.satisfies(&r_exact));
474
475        assert!(r3.satisfies(&r1));
476        assert!(r4.satisfies(&r1));
477        assert!(!r1.satisfies(&r3));
478        assert!(!r1.satisfies(&r4));
479        assert!(!r3.satisfies(&r4));
480        assert!(!r4.satisfies(&r3));
481
482        assert!(r_exact.satisfies(&r1));
483        assert!(!r1.satisfies(&r_exact));
484        assert!(!r3.satisfies(&r_exact));
485        assert!(!r_exact.satisfies(&r3));
486
487        assert!(r3.satisfies(&r_exact_single));
488        assert!(r_exact_single.satisfies(&r3));
489    }
490}
risingwave_frontend/optimizer/property/distribution.rs

risingwave_frontend/optimizer/property/
distribution.rs
