Merge pull request #11823 from aalexandrov/protobuf_variadic_func

MaterializeInc · Apr 19, 2022 · 2d3e35a · 2d3e35a
2 parents a69372b + b9a0f89
commit 2d3e35a
Show file tree

Hide file tree

Showing 2 changed files with 151 additions and 0 deletions.
diff --git a/src/expr/src/scalar/func.proto b/src/expr/src/scalar/func.proto
@@ -520,3 +520,37 @@ message ProtoBinaryFunc {
         google.protobuf.Empty power_numeric = 139;
     }
 }
+
+message ProtoVariadicFunc {
+    oneof kind {
+        google.protobuf.Empty coalesce = 1;
+        google.protobuf.Empty greatest = 2;
+        google.protobuf.Empty least = 3;
+        google.protobuf.Empty concat = 4;
+        google.protobuf.Empty make_timestamp = 5;
+        google.protobuf.Empty pad_leading = 6;
+        google.protobuf.Empty substr = 7;
+        google.protobuf.Empty replace = 8;
+        google.protobuf.Empty jsonb_build_array = 9;
+        google.protobuf.Empty jsonb_build_object = 10;
+        // unsupported: ArrayCreate { elem_type: ScalarType }
+        google.protobuf.Empty array_create = 11;
+        // unsupported: ArrayToString { elem_type: ScalarType }
+        google.protobuf.Empty array_to_string = 12;
+        // unsupported: ArrayIndex { offset: usize }
+        google.protobuf.Empty array_index = 13;
+        // unsupported: ListCreate { elem_type: ScalarType }
+        google.protobuf.Empty list_create = 14;
+        // unsupported: RecordCreate { field_names: Vec<ColumnName> }
+        google.protobuf.Empty record_create = 15;
+        google.protobuf.Empty list_index = 16;
+        google.protobuf.Empty list_slice_linear = 17;
+        google.protobuf.Empty split_part = 18;
+        google.protobuf.Empty regexp_match = 19;
+        google.protobuf.Empty hmac_string = 20;
+        google.protobuf.Empty hmac_bytes = 21;
+        google.protobuf.Empty error_if_null = 22;
+        google.protobuf.Empty date_bin_timestamp = 23;
+        google.protobuf.Empty date_bin_timestamp_tz = 24;
+    }
+}
diff --git a/src/expr/src/scalar/func.rs b/src/expr/src/scalar/func.rs
@@ -7144,6 +7144,116 @@ impl fmt::Display for VariadicFunc {
     }
 }
 
+impl Arbitrary for VariadicFunc {
+    type Parameters = ();
+
+    type Strategy = Union<BoxedStrategy<Self>>;
+
+    fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
+        prop_oneof![
+            Just(VariadicFunc::Coalesce),
+            Just(VariadicFunc::Greatest),
+            Just(VariadicFunc::Least),
+            Just(VariadicFunc::Concat),
+            Just(VariadicFunc::MakeTimestamp),
+            Just(VariadicFunc::PadLeading),
+            Just(VariadicFunc::Substr),
+            Just(VariadicFunc::Replace),
+            Just(VariadicFunc::JsonbBuildArray),
+            Just(VariadicFunc::JsonbBuildObject),
+            // todo: ArrayCreate { elem_type: ScalarType },
+            // todo: ArrayToString { elem_type: ScalarType },
+            // todo: ArrayIndex { offset: usize },
+            // todo: ListCreate { elem_type: ScalarType },
+            // todo: RecordCreate { field_names: Vec<ColumnName> },
+            Just(VariadicFunc::ListIndex),
+            Just(VariadicFunc::ListSliceLinear),
+            Just(VariadicFunc::SplitPart),
+            Just(VariadicFunc::RegexpMatch),
+            Just(VariadicFunc::HmacString),
+            Just(VariadicFunc::HmacBytes),
+            Just(VariadicFunc::ErrorIfNull),
+            Just(VariadicFunc::DateBinTimestamp),
+            Just(VariadicFunc::DateBinTimestampTz),
+        ]
+    }
+}
+
+impl From<&VariadicFunc> for ProtoVariadicFunc {
+    #[allow(clippy::todo)]
+    fn from(func: &VariadicFunc) -> Self {
+        use proto_variadic_func::Kind::*;
+        let kind = match func {
+            VariadicFunc::Coalesce => Coalesce(()),
+            VariadicFunc::Greatest => Greatest(()),
+            VariadicFunc::Least => Least(()),
+            VariadicFunc::Concat => Concat(()),
+            VariadicFunc::MakeTimestamp => MakeTimestamp(()),
+            VariadicFunc::PadLeading => PadLeading(()),
+            VariadicFunc::Substr => Substr(()),
+            VariadicFunc::Replace => Replace(()),
+            VariadicFunc::JsonbBuildArray => JsonbBuildArray(()),
+            VariadicFunc::JsonbBuildObject => JsonbBuildObject(()),
+            VariadicFunc::ArrayCreate { .. } => todo!(),
+            VariadicFunc::ArrayToString { .. } => todo!(),
+            VariadicFunc::ArrayIndex { .. } => todo!(),
+            VariadicFunc::ListCreate { .. } => todo!(),
+            VariadicFunc::RecordCreate { .. } => todo!(),
+            VariadicFunc::ListIndex => ListIndex(()),
+            VariadicFunc::ListSliceLinear => ListSliceLinear(()),
+            VariadicFunc::SplitPart => SplitPart(()),
+            VariadicFunc::RegexpMatch => RegexpMatch(()),
+            VariadicFunc::HmacString => HmacString(()),
+            VariadicFunc::HmacBytes => HmacBytes(()),
+            VariadicFunc::ErrorIfNull => ErrorIfNull(()),
+            VariadicFunc::DateBinTimestamp => DateBinTimestamp(()),
+            VariadicFunc::DateBinTimestampTz => DateBinTimestampTz(()),
+        };
+        ProtoVariadicFunc { kind: Some(kind) }
+    }
+}
+
+impl TryFrom<ProtoVariadicFunc> for VariadicFunc {
+    type Error = TryFromProtoError;
+
+    #[allow(clippy::todo)]
+    fn try_from(func: ProtoVariadicFunc) -> Result<Self, Self::Error> {
+        use proto_variadic_func::Kind::*;
+        if let Some(kind) = func.kind {
+            match kind {
+                Coalesce(()) => Ok(VariadicFunc::Coalesce),
+                Greatest(()) => Ok(VariadicFunc::Greatest),
+                Least(()) => Ok(VariadicFunc::Least),
+                Concat(()) => Ok(VariadicFunc::Concat),
+                MakeTimestamp(()) => Ok(VariadicFunc::MakeTimestamp),
+                PadLeading(()) => Ok(VariadicFunc::PadLeading),
+                Substr(()) => Ok(VariadicFunc::Substr),
+                Replace(()) => Ok(VariadicFunc::Replace),
+                JsonbBuildArray(()) => Ok(VariadicFunc::JsonbBuildArray),
+                JsonbBuildObject(()) => Ok(VariadicFunc::JsonbBuildObject),
+                ArrayCreate(()) => todo!(),
+                ArrayToString(()) => todo!(),
+                ArrayIndex(()) => todo!(),
+                ListCreate(()) => todo!(),
+                RecordCreate(()) => todo!(),
+                ListIndex(()) => Ok(VariadicFunc::ListIndex),
+                ListSliceLinear(()) => Ok(VariadicFunc::ListSliceLinear),
+                SplitPart(()) => Ok(VariadicFunc::SplitPart),
+                RegexpMatch(()) => Ok(VariadicFunc::RegexpMatch),
+                HmacString(()) => Ok(VariadicFunc::HmacString),
+                HmacBytes(()) => Ok(VariadicFunc::HmacBytes),
+                ErrorIfNull(()) => Ok(VariadicFunc::ErrorIfNull),
+                DateBinTimestamp(()) => Ok(VariadicFunc::DateBinTimestamp),
+                DateBinTimestampTz(()) => Ok(VariadicFunc::DateBinTimestampTz),
+            }
+        } else {
+            Err(TryFromProtoError::missing_field(
+                "`ProtoVariadicFunc::kind`",
+            ))
+        }
+    }
+}
+
 #[cfg(test)]
 mod test {
     use chrono::prelude::*;
@@ -7258,5 +7368,12 @@ mod test {
             assert!(actual.is_ok());
             assert_eq!(actual.unwrap(), expect);
         }
+
+        #[test]
+        fn variadic_func_protobuf_roundtrip(expect in any::<VariadicFunc>()) {
+            let actual = protobuf_roundtrip::<_, ProtoVariadicFunc>(&expect);
+            assert!(actual.is_ok());
+            assert_eq!(actual.unwrap(), expect);
+        }
     }
 }