Skip to main content

datafusion_expr/
proto.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! Protobuf conversions for the expression types owned by this crate:
19//! [`WindowFrame`], [`WindowFrameBound`], [`WindowFrameUnits`],
20//! [`MergeIntoClauseKind`](crate::dml::MergeIntoClauseKind) and
21//! [`NullTreatment`](crate::expr::NullTreatment).
22//!
23//! These are plain [`From`] / [`TryFrom`] impls rather than something taking a
24//! codec: every field is either an enum tag or a [`ScalarValue`], so the
25//! conversion needs nothing but the value itself. The orphan rule allows them
26//! here because one side of each conversion is a type this crate owns.
27//!
28//! [`ScalarValue`]: datafusion_common::ScalarValue
29
30use datafusion_common::ScalarValue;
31use datafusion_proto_common::{FromProtoError, ToProtoError};
32use datafusion_proto_models::protobuf;
33
34use crate::dml::MergeIntoClauseKind;
35use crate::expr::NullTreatment;
36use crate::{WindowFrame, WindowFrameBound, WindowFrameUnits};
37
38impl From<protobuf::WindowFrameUnits> for WindowFrameUnits {
39    fn from(units: protobuf::WindowFrameUnits) -> Self {
40        match units {
41            protobuf::WindowFrameUnits::Rows => Self::Rows,
42            protobuf::WindowFrameUnits::Range => Self::Range,
43            protobuf::WindowFrameUnits::Groups => Self::Groups,
44        }
45    }
46}
47
48impl From<WindowFrameUnits> for protobuf::WindowFrameUnits {
49    fn from(units: WindowFrameUnits) -> Self {
50        match units {
51            WindowFrameUnits::Rows => Self::Rows,
52            WindowFrameUnits::Range => Self::Range,
53            WindowFrameUnits::Groups => Self::Groups,
54        }
55    }
56}
57
58impl TryFrom<protobuf::WindowFrameBound> for WindowFrameBound {
59    type Error = FromProtoError;
60
61    fn try_from(bound: protobuf::WindowFrameBound) -> Result<Self, Self::Error> {
62        let bound_type =
63            protobuf::WindowFrameBoundType::try_from(bound.window_frame_bound_type)
64                .map_err(|_| {
65                    FromProtoError::unknown(
66                        "WindowFrameBoundType",
67                        bound.window_frame_bound_type,
68                    )
69                })?;
70        match bound_type {
71            protobuf::WindowFrameBoundType::CurrentRow => Ok(Self::CurrentRow),
72            protobuf::WindowFrameBoundType::Preceding => match bound.bound_value {
73                Some(x) => Ok(Self::Preceding(ScalarValue::try_from(&x)?)),
74                None => Ok(Self::Preceding(ScalarValue::UInt64(None))),
75            },
76            protobuf::WindowFrameBoundType::Following => match bound.bound_value {
77                Some(x) => Ok(Self::Following(ScalarValue::try_from(&x)?)),
78                None => Ok(Self::Following(ScalarValue::UInt64(None))),
79            },
80        }
81    }
82}
83
84impl TryFrom<&WindowFrameBound> for protobuf::WindowFrameBound {
85    type Error = ToProtoError;
86
87    fn try_from(bound: &WindowFrameBound) -> Result<Self, Self::Error> {
88        Ok(match bound {
89            WindowFrameBound::CurrentRow => Self {
90                window_frame_bound_type: protobuf::WindowFrameBoundType::CurrentRow
91                    .into(),
92                bound_value: None,
93            },
94            WindowFrameBound::Preceding(v) => Self {
95                window_frame_bound_type: protobuf::WindowFrameBoundType::Preceding.into(),
96                bound_value: Some(v.try_into()?),
97            },
98            WindowFrameBound::Following(v) => Self {
99                window_frame_bound_type: protobuf::WindowFrameBoundType::Following.into(),
100                bound_value: Some(v.try_into()?),
101            },
102        })
103    }
104}
105
106impl TryFrom<protobuf::WindowFrame> for WindowFrame {
107    type Error = FromProtoError;
108
109    fn try_from(window: protobuf::WindowFrame) -> Result<Self, Self::Error> {
110        let units = WindowFrameUnits::from(
111            protobuf::WindowFrameUnits::try_from(window.window_frame_units).map_err(
112                |_| {
113                    FromProtoError::unknown("WindowFrameUnits", window.window_frame_units)
114                },
115            )?,
116        );
117        let start_bound = WindowFrameBound::try_from(
118            window
119                .start_bound
120                .ok_or_else(|| FromProtoError::required("start_bound"))?,
121        )?;
122        let end_bound = window
123            .end_bound
124            .map(|end_bound| match end_bound {
125                protobuf::window_frame::EndBound::Bound(end_bound) => {
126                    WindowFrameBound::try_from(end_bound)
127                }
128            })
129            .transpose()?
130            .unwrap_or(WindowFrameBound::CurrentRow);
131        Ok(WindowFrame::new_bounds(units, start_bound, end_bound))
132    }
133}
134
135impl TryFrom<&WindowFrame> for protobuf::WindowFrame {
136    type Error = ToProtoError;
137
138    fn try_from(window: &WindowFrame) -> Result<Self, Self::Error> {
139        Ok(Self {
140            window_frame_units: protobuf::WindowFrameUnits::from(window.units).into(),
141            start_bound: Some((&window.start_bound).try_into()?),
142            end_bound: Some(protobuf::window_frame::EndBound::Bound(
143                (&window.end_bound).try_into()?,
144            )),
145        })
146    }
147}
148
149impl From<protobuf::merge_into_clause_node::Kind> for MergeIntoClauseKind {
150    fn from(kind: protobuf::merge_into_clause_node::Kind) -> Self {
151        match kind {
152            protobuf::merge_into_clause_node::Kind::Matched => Self::Matched,
153            protobuf::merge_into_clause_node::Kind::NotMatched => Self::NotMatched,
154            protobuf::merge_into_clause_node::Kind::NotMatchedByTarget => {
155                Self::NotMatchedByTarget
156            }
157            protobuf::merge_into_clause_node::Kind::NotMatchedBySource => {
158                Self::NotMatchedBySource
159            }
160        }
161    }
162}
163
164impl From<MergeIntoClauseKind> for protobuf::merge_into_clause_node::Kind {
165    fn from(kind: MergeIntoClauseKind) -> Self {
166        match kind {
167            MergeIntoClauseKind::Matched => Self::Matched,
168            MergeIntoClauseKind::NotMatched => Self::NotMatched,
169            MergeIntoClauseKind::NotMatchedByTarget => Self::NotMatchedByTarget,
170            MergeIntoClauseKind::NotMatchedBySource => Self::NotMatchedBySource,
171        }
172    }
173}
174
175impl From<protobuf::NullTreatment> for NullTreatment {
176    fn from(t: protobuf::NullTreatment) -> Self {
177        match t {
178            protobuf::NullTreatment::RespectNulls => Self::RespectNulls,
179            protobuf::NullTreatment::IgnoreNulls => Self::IgnoreNulls,
180        }
181    }
182}
183
184impl From<NullTreatment> for protobuf::NullTreatment {
185    fn from(t: NullTreatment) -> Self {
186        match t {
187            NullTreatment::RespectNulls => Self::RespectNulls,
188            NullTreatment::IgnoreNulls => Self::IgnoreNulls,
189        }
190    }
191}
192
193#[cfg(test)]
194mod tests {
195    use super::*;
196
197    #[test]
198    fn window_frame_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
199        let frame = WindowFrame::new_bounds(
200            WindowFrameUnits::Range,
201            WindowFrameBound::Preceding(ScalarValue::UInt64(Some(2))),
202            WindowFrameBound::Following(ScalarValue::UInt64(Some(3))),
203        );
204
205        let encoded = protobuf::WindowFrame::try_from(&frame)?;
206        let decoded = WindowFrame::try_from(encoded)?;
207
208        assert_eq!(decoded.units, frame.units);
209        assert_eq!(decoded.start_bound, frame.start_bound);
210        assert_eq!(decoded.end_bound, frame.end_bound);
211        Ok(())
212    }
213
214    #[test]
215    fn window_frame_from_proto_rejects_missing_start_bound() {
216        let proto = protobuf::WindowFrame {
217            window_frame_units: protobuf::WindowFrameUnits::Rows.into(),
218            start_bound: None,
219            end_bound: None,
220        };
221
222        let err = WindowFrame::try_from(proto).unwrap_err();
223        assert!(
224            err.to_string().contains("start_bound"),
225            "unexpected error: {err}"
226        );
227    }
228
229    #[test]
230    fn missing_end_bound_decodes_as_current_row() -> Result<(), Box<dyn std::error::Error>>
231    {
232        let proto = protobuf::WindowFrame {
233            window_frame_units: protobuf::WindowFrameUnits::Rows.into(),
234            start_bound: Some(protobuf::WindowFrameBound {
235                window_frame_bound_type: protobuf::WindowFrameBoundType::CurrentRow
236                    .into(),
237                bound_value: None,
238            }),
239            end_bound: None,
240        };
241
242        let decoded = WindowFrame::try_from(proto)?;
243        assert_eq!(decoded.end_bound, WindowFrameBound::CurrentRow);
244        Ok(())
245    }
246}