operator/write_dispatcher/
attribution.rs

1// Copyright 2023 Greptime Team
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
15use std::sync::Arc;
16
17use common_meta::error::{Result, UnexpectedSnafu};
18use partition::partition::RegionMask;
19use session::write_progress::WriteProgress;
20use snafu::ensure;
21
22/// Row ownership in payload order. Kept in process memory, never encoded in an RPC.
23#[derive(Clone, Debug, Default)]
24pub struct WriteAttribution {
25    spans: Vec<(usize, Option<Arc<WriteProgress>>)>,
26}
27
28impl WriteAttribution {
29    pub fn new(rows: usize, progress: Option<Arc<WriteProgress>>) -> Self {
30        Self {
31            spans: vec![(rows, progress)],
32        }
33    }
34
35    pub fn extend(&mut self, other: Self) {
36        self.spans.extend(other.spans);
37    }
38
39    pub fn row_count(&self) -> usize {
40        self.spans.iter().map(|(rows, _)| rows).sum()
41    }
42
43    pub fn is_tracked(&self) -> bool {
44        self.spans.iter().any(|(_, progress)| progress.is_some())
45    }
46
47    /// Selects ownership using the exact mask applied to the Arrow payload.
48    pub fn filter(&self, mask: &RegionMask) -> Result<Self> {
49        ensure!(
50            self.row_count() == mask.array().len(),
51            UnexpectedSnafu {
52                err_msg: "Write ownership does not match the partition mask",
53            }
54        );
55        if !self.is_tracked() {
56            return Ok(Self::default());
57        }
58        let mut offset = 0;
59        let mut selected = Self::default();
60        for (rows, progress) in &self.spans {
61            let count = mask.array().slice(offset, *rows).true_count();
62            if count > 0 {
63                selected.spans.push((count, progress.clone()));
64            }
65            offset += rows;
66        }
67        Ok(selected)
68    }
69
70    /// Successful writes must acknowledge every input row, including duplicates, as
71    /// Mito's `RegionWriteCtx::push_bulk` does. A partial count cannot identify owners.
72    pub(crate) fn record(&self, affected_rows: usize) -> Result<()> {
73        if !self.is_tracked() {
74            return Ok(());
75        }
76        ensure!(
77            affected_rows == self.row_count(),
78            UnexpectedSnafu {
79                err_msg: format!(
80                    "Write affected {affected_rows} rows, expected {}; cannot attribute the result",
81                    self.row_count()
82                ),
83            }
84        );
85        for (rows, progress) in &self.spans {
86            if let Some(progress) = progress {
87                progress.add_written(*rows);
88            }
89        }
90        Ok(())
91    }
92
93    /// Direct row writes have one owner even when a refreshed route splits them.
94    pub(crate) fn for_rows(&self, rows: usize) -> Result<Self> {
95        ensure!(
96            self.spans.len() <= 1,
97            UnexpectedSnafu {
98                err_msg: "Cannot repartition shared ownership without a row mask",
99            }
100        );
101        Ok(Self::new(
102            rows,
103            self.spans.first().and_then(|(_, p)| p.clone()),
104        ))
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use datatypes::arrow::array::BooleanArray;
111
112    use super::*;
113
114    #[test]
115    fn ownership_survives_filtering_and_repartition() {
116        let first = Arc::new(WriteProgress::default());
117        let second = Arc::new(WriteProgress::default());
118        let mut owners = WriteAttribution::new(3, Some(first.clone()));
119        owners.extend(WriteAttribution::new(2, Some(second.clone())));
120        let initial = owners
121            .filter(&RegionMask::new(
122                BooleanArray::from(vec![true, false, true, false, true]),
123                3,
124            ))
125            .unwrap();
126        let retry = initial
127            .filter(&RegionMask::new(
128                BooleanArray::from(vec![false, true, true]),
129                2,
130            ))
131            .unwrap();
132        assert!(retry.record(1).is_err());
133        assert_eq!(first.written(), 0);
134        retry.record(2).unwrap();
135        assert_eq!(first.written(), 1);
136        assert_eq!(second.written(), 1);
137        assert!(owners.for_rows(5).is_err());
138        assert!(
139            owners
140                .filter(&RegionMask::new(BooleanArray::from(vec![true]), 1))
141                .is_err()
142        );
143    }
144}
operator/write_dispatcher/attribution.rs

operator/write_dispatcher/
attribution.rs
