risingwave_connector/sink/formatter/
upsert.rs

1// Copyright 2025 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
15use risingwave_common::array::Op;
16
17use super::{Result, SinkFormatter, StreamChunk};
18use crate::sink::encoder::RowEncoder;
19use crate::tri;
20
21pub struct UpsertFormatter<KE, VE> {
22    key_encoder: KE,
23    val_encoder: VE,
24}
25
26impl<KE, VE> UpsertFormatter<KE, VE> {
27    pub fn new(key_encoder: KE, val_encoder: VE) -> Self {
28        Self {
29            key_encoder,
30            val_encoder,
31        }
32    }
33}
34
35impl<KE: RowEncoder, VE: RowEncoder> SinkFormatter for UpsertFormatter<KE, VE> {
36    type K = KE::Output;
37    type V = VE::Output;
38
39    fn format_chunk(
40        &self,
41        chunk: &StreamChunk,
42    ) -> impl Iterator<Item = Result<(Option<Self::K>, Option<Self::V>)>> {
43        std::iter::from_coroutine(
44            #[coroutine]
45            || {
46                for (op, row) in chunk.rows() {
47                    let event_key_object = Some(tri!(self.key_encoder.encode(row)));
48
49                    let event_object = match op {
50                        Op::Insert | Op::UpdateInsert => Some(tri!(self.val_encoder.encode(row))),
51                        // Empty value with a key
52                        Op::Delete => None,
53                        Op::UpdateDelete => {
54                            // upsert semantic does not require update delete event
55                            continue;
56                        }
57                    };
58
59                    yield Ok((event_key_object, event_object))
60                }
61            },
62        )
63    }
64}
risingwave_connector/sink/formatter/upsert.rs

risingwave_connector/sink/formatter/
upsert.rs
