make starts_with and ends_with work with bytes https://github.com/cozodb/cozo/discussions/132

cozo-core/src/data/functions.rs

@@ -1359,28 +1359,24 @@ pub(crate) fn op_trim_end(args: &[DataValue]) -> Result<DataValue> {
 
 define_op!(OP_STARTS_WITH, 2, false);
 pub(crate) fn op_starts_with(args: &[DataValue]) -> Result<DataValue> {
-    let a = match &args[0] {
-        DataValue::Str(s) => s,
-        _ => bail!("'starts_with' requires strings"),
-    };
-    let b = match &args[1] {
-        DataValue::Str(s) => s,
-        _ => bail!("'starts_with' requires strings"),
-    };
-    Ok(DataValue::from(a.starts_with(b as &str)))
+    match (&args[0], &args[1]) {
+        (DataValue::Str(l), DataValue::Str(r)) => Ok(DataValue::from(l.starts_with(r as &str))),
+        (DataValue::Bytes(l), DataValue::Bytes(r)) => {
+            Ok(DataValue::from(l.starts_with(r as &[u8])))
+        }
+        _ => bail!("'starts_with' requires strings or bytes"),
+    }
 }
 
 define_op!(OP_ENDS_WITH, 2, false);
 pub(crate) fn op_ends_with(args: &[DataValue]) -> Result<DataValue> {
-    let a = match &args[0] {
-        DataValue::Str(s) => s,
-        _ => bail!("'ends_with' requires strings"),
-    };
-    let b = match &args[1] {
-        DataValue::Str(s) => s,
-        _ => bail!("'ends_with' requires strings"),
-    };
-    Ok(DataValue::from(a.ends_with(b as &str)))
+    match (&args[0], &args[1]) {
+        (DataValue::Str(l), DataValue::Str(r)) => Ok(DataValue::from(l.ends_with(r as &str))),
+        (DataValue::Bytes(l), DataValue::Bytes(r)) => {
+            Ok(DataValue::from(l.ends_with(r as &[u8])))
+        }
+        _ => bail!("'ends_with' requires strings or bytes"),
+    }
 }
 
 define_op!(OP_REGEX, 1, false);
@@ -2086,16 +2082,28 @@ pub(crate) fn op_vec(args: &[DataValue]) -> Result<DataValue> {
             }
         },
         DataValue::Str(s) => {
-            let bytes = STANDARD.decode(s).map_err(|_| miette!("Data is not base64 encoded"))?;
+            let bytes = STANDARD
+                .decode(s)
+                .map_err(|_| miette!("Data is not base64 encoded"))?;
             match t {
                 VecElementType::F32 => {
                     let f32_count = bytes.len() / mem::size_of::<f32>();
-                    let arr = unsafe { ndarray::ArrayView1::from_shape_ptr(ndarray::Dim([f32_count]), bytes.as_ptr() as *const f32) };
+                    let arr = unsafe {
+                        ndarray::ArrayView1::from_shape_ptr(
+                            ndarray::Dim([f32_count]),
+                            bytes.as_ptr() as *const f32,
+                        )
+                    };
                     Ok(DataValue::Vec(Vector::F32(arr.to_owned())))
                 }
                 VecElementType::F64 => {
                     let f64_count = bytes.len() / mem::size_of::<f64>();
-                    let arr = unsafe { ndarray::ArrayView1::from_shape_ptr(ndarray::Dim([f64_count]), bytes.as_ptr() as *const f64) };
+                    let arr = unsafe {
+                        ndarray::ArrayView1::from_shape_ptr(
+                            ndarray::Dim([f64_count]),
+                            bytes.as_ptr() as *const f64,
+                        )
+                    };
                     Ok(DataValue::Vec(Vector::F64(arr.to_owned())))
                 }
             }