perf(expr): support writer-style #[function] for bytea type (#23727)

Signed-off-by: Mingzhuo Yin <[email protected]>

Author: silver-ymz

Cargo.lock

@@ -12,8 +12,9 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::io::Write;
 use std::iter;
-use std::mem::size_of;
+use std::mem::{ManuallyDrop, size_of};
 
 use risingwave_common_estimate_size::EstimateSize;
 use risingwave_pb::common::Buffer;
@@ -250,33 +251,13 @@ impl BytesArrayBuilder {
     }
 }
 
+/// Note: dropping an unfinished `BytesWriter` will rollback the partial data
 pub struct BytesWriter<'a> {
     builder: &'a mut BytesArrayBuilder,
 }
 
-impl<'a> BytesWriter<'a> {
-    /// `write_ref` will consume `BytesWriter` and pass the ownership of `builder` to `BytesGuard`.
-    pub fn write_ref(self, value: &[u8]) {
-        self.builder.append(Some(value));
-    }
-
-    /// `begin` will create a `PartialBytesWriter`, which allow multiple appendings to create a new
-    /// record.
-    pub fn begin(self) -> PartialBytesWriter<'a> {
-        PartialBytesWriter {
-            builder: self.builder,
-        }
-    }
-}
-
-pub struct PartialBytesWriter<'a> {
-    builder: &'a mut BytesArrayBuilder,
-}
-
-impl PartialBytesWriter<'_> {
+impl BytesWriter<'_> {
     /// `write_ref` will append partial dirty data to `builder`.
-    /// `PartialBytesWriter::write_ref` is different from `BytesWriter::write_ref`
-    /// in that it allows us to call it multiple times.
     pub fn write_ref(&mut self, value: &[u8]) {
         // SAFETY: We'll clean the dirty `builder` in the `drop`.
         unsafe { self.builder.append_partial(value) }
@@ -286,10 +267,27 @@ impl PartialBytesWriter<'_> {
     /// Exactly one new record was appended and the `builder` can be safely used.
     pub fn finish(self) {
         self.builder.finish_partial();
+        let _ = ManuallyDrop::new(self); // Prevent drop
+    }
+}
+
+impl Write for BytesWriter<'_> {
+    fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+        self.write_ref(buf);
+        Ok(buf.len())
+    }
+
+    fn write_all(&mut self, buf: &[u8]) -> std::io::Result<()> {
+        self.write_ref(buf);
+        Ok(())
+    }
+
+    fn flush(&mut self) -> std::io::Result<()> {
+        Ok(())
     }
 }
 
-impl Drop for PartialBytesWriter<'_> {
+impl Drop for BytesWriter<'_> {
     fn drop(&mut self) {
         // If `finish` is not called, we should rollback the data.
         self.builder.rollback_partial();

src/common/src/array/bytes_array.rs

@@ -17,7 +17,7 @@ use std::fmt::{Display, Write};
 use risingwave_common_estimate_size::EstimateSize;
 use risingwave_pb::data::{ArrayType, PbArray};
 
-use super::bytes_array::{BytesWriter, PartialBytesWriter};
+use super::bytes_array::BytesWriter;
 use super::{Array, ArrayBuilder, BytesArray, BytesArrayBuilder, DataType};
 use crate::bitmap::Bitmap;
 
@@ -105,7 +105,7 @@ impl Utf8Array {
         let mut builder = Utf8ArrayBuilder::new(iter.size_hint().0);
         for e in iter {
             if let Some(s) = e {
-                let mut writer = builder.writer().begin();
+                let mut writer = builder.writer();
                 write!(writer, "{}", s).unwrap();
                 writer.finish();
             } else {
@@ -177,7 +177,7 @@ impl Utf8ArrayBuilder {
     /// Append an element as the `Display` format to the array.
     pub fn append_display(&mut self, value: Option<impl Display>) {
         if let Some(s) = value {
-            let mut writer = self.writer().begin();
+            let mut writer = self.writer();
             write!(writer, "{}", s).unwrap();
             writer.finish();
         } else {
@@ -186,35 +186,20 @@ impl Utf8ArrayBuilder {
     }
 }
 
+/// Note: dropping an unfinished `StringWriter` will rollback the partial data, which is the behavior of the inner `BytesWriter`.
 pub struct StringWriter<'a> {
     bytes: BytesWriter<'a>,
 }
 
-impl<'a> StringWriter<'a> {
-    /// `begin` will create a `PartialStringWriter`, which allow multiple appendings to create a new
-    /// record.
-    pub fn begin(self) -> PartialStringWriter<'a> {
-        PartialStringWriter {
-            bytes: self.bytes.begin(),
-        }
-    }
-}
-
-// Note: dropping an unfinished `PartialStringWriter` will rollback the partial data, which is the
-// behavior of the inner `PartialBytesWriter`.
-pub struct PartialStringWriter<'a> {
-    bytes: PartialBytesWriter<'a>,
-}
-
-impl PartialStringWriter<'_> {
+impl StringWriter<'_> {
     /// `finish` will be called while the entire record is written.
     /// Exactly one new record was appended and the `builder` can be safely used.
     pub fn finish(self) {
         self.bytes.finish()
     }
 }
 
-impl Write for PartialStringWriter<'_> {
+impl Write for StringWriter<'_> {
     fn write_str(&mut self, s: &str) -> std::fmt::Result {
         self.bytes.write_ref(s.as_bytes());
         Ok(())
@@ -245,15 +230,14 @@ mod tests {
     }
 
     #[test]
-    fn test_utf8_partial_writer() {
+    fn test_utf8_writer() {
         let mut builder = Utf8ArrayBuilder::new(0);
         {
-            let writer = builder.writer();
-            let mut partial_writer = writer.begin();
+            let mut writer = builder.writer();
             for _ in 0..2 {
-                partial_writer.write_str("ran").unwrap();
+                writer.write_str("ran").unwrap();
             }
-            partial_writer.finish()
+            writer.finish()
         };
         let array = builder.finish();
         assert_eq!(array.len(), 1);
@@ -262,33 +246,30 @@ mod tests {
     }
 
     #[test]
-    fn test_utf8_partial_writer_failed() {
+    fn test_utf8_writer_failed() {
         let mut builder = Utf8ArrayBuilder::new(0);
         // Write a record.
         {
-            let writer = builder.writer();
-            let mut partial_writer = writer.begin();
-            partial_writer.write_str("Dia").unwrap();
-            partial_writer.write_str("na").unwrap();
-            partial_writer.finish()
+            let mut writer = builder.writer();
+            writer.write_str("Dia").unwrap();
+            writer.write_str("na").unwrap();
+            writer.finish()
         };
 
         // Write a record failed.
         {
-            let writer = builder.writer();
-            let mut partial_writer = writer.begin();
-            partial_writer.write_str("Ca").unwrap();
-            partial_writer.write_str("rol").unwrap();
+            let mut writer = builder.writer();
+            writer.write_str("Ca").unwrap();
+            writer.write_str("rol").unwrap();
             // We don't finish here.
         };
 
         // Write a record.
         {
-            let writer = builder.writer();
-            let mut partial_writer = writer.begin();
-            partial_writer.write_str("Ki").unwrap();
-            partial_writer.write_str("ra").unwrap();
-            partial_writer.finish()
+            let mut writer = builder.writer();
+            writer.write_str("Ki").unwrap();
+            writer.write_str("ra").unwrap();
+            writer.finish()
         };
 
         // Verify only two valid records.

src/common/src/array/utf8_array.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 use itertools::Itertools;
+use thiserror_ext::AsReport;
 
 use crate::types::{Timestamp, Timestamptz};
 
@@ -113,11 +114,11 @@ pub fn i64_to_timestamp(t: i64) -> Result<Timestamp> {
 }
 
 /// Refer to PostgreSQL's implementation <https://github.com/postgres/postgres/blob/5cb54fc310fb84287cbdc74533f3420490a2f63a/src/backend/utils/adt/varlena.c#L276-L288>
-pub fn str_to_bytea(elem: &str) -> Result<Box<[u8]>> {
+pub fn str_to_bytea(elem: &str, writer: &mut impl std::io::Write) -> Result<()> {
     if let Some(remainder) = elem.strip_prefix(r"\x") {
-        Ok(parse_bytes_hex(remainder)?.into())
+        Ok(parse_bytes_hex(remainder, writer)?)
     } else {
-        Ok(parse_bytes_traditional(elem)?.into())
+        Ok(parse_bytes_traditional(elem, writer)?)
     }
 }
 
@@ -132,9 +133,7 @@ fn get_hex(c: u8) -> Result<u8> {
 }
 
 /// Refer to <https://www.postgresql.org/docs/current/datatype-binary.html#id-1.5.7.12.10> for specification.
-pub fn parse_bytes_hex(s: &str) -> Result<Vec<u8>> {
-    let mut res = Vec::with_capacity(s.len() / 2);
-
+pub fn parse_bytes_hex(s: &str, writer: &mut impl std::io::Write) -> Result<()> {
     let mut bytes = s.bytes();
     while let Some(c) = bytes.next() {
         // white spaces are tolerated
@@ -146,31 +145,35 @@ pub fn parse_bytes_hex(s: &str) -> Result<Vec<u8>> {
         match bytes.next() {
             Some(c) => {
                 let v2 = get_hex(c)?;
-                res.push((v1 << 4) | v2);
+                writer
+                    .write_all(&[(v1 << 4) | v2])
+                    .map_err(|e| e.to_report_string())?;
             }
             None => return Err("invalid hexadecimal data: odd number of digits".to_owned()),
         }
     }
 
-    Ok(res)
+    Ok(())
 }
 
 /// Refer to <https://www.postgresql.org/docs/current/datatype-binary.html#id-1.5.7.12.10> for specification.
-pub fn parse_bytes_traditional(s: &str) -> Result<Vec<u8>> {
+pub fn parse_bytes_traditional(s: &str, writer: &mut impl std::io::Write) -> Result<()> {
     let mut bytes = s.bytes();
 
-    let mut res = Vec::new();
     while let Some(b) = bytes.next() {
         if b != b'\\' {
-            res.push(b);
+            writer.write_all(&[b]).map_err(|e| e.to_report_string())?;
         } else {
             match bytes.next() {
                 Some(b'\\') => {
-                    res.push(b'\\');
+                    writer.write_all(b"\\").map_err(|e| e.to_report_string())?;
                 }
                 Some(b1 @ b'0'..=b'3') => match bytes.next_tuple() {
                     Some((b2 @ b'0'..=b'7', b3 @ b'0'..=b'7')) => {
-                        res.push(((b1 - b'0') << 6) + ((b2 - b'0') << 3) + (b3 - b'0'));
+                        let byte = (b1 - b'0') << 6 | (b2 - b'0') << 3 | (b3 - b'0');
+                        writer
+                            .write_all(&[byte])
+                            .map_err(|e| e.to_report_string())?;
                     }
                     _ => {
                         // one backslash, not followed by another or ### valid octal
@@ -185,7 +188,7 @@ pub fn parse_bytes_traditional(s: &str) -> Result<Vec<u8>> {
         }
     }
 
-    Ok(res)
+    Ok(())
 }
 
 #[cfg(test)]
@@ -206,6 +209,13 @@ mod tests {
     #[test]
     fn test_bytea() {
         use crate::types::ToText;
+
+        fn str_to_bytea(s: &str) -> Result<Box<[u8]>> {
+            let mut buf = Vec::new();
+            super::str_to_bytea(s, &mut buf)?;
+            Ok(buf.into())
+        }
+
         assert_eq!(str_to_bytea("fgo").unwrap().as_ref().to_text(), r"\x66676f");
         assert_eq!(
             str_to_bytea(r"\xDeadBeef").unwrap().as_ref().to_text(),

src/common/src/cast/mod.rs

@@ -1091,7 +1091,11 @@ impl ScalarImpl {
             DataType::List(_) => ListValue::from_str(s, data_type)?.into(),
             DataType::Struct(st) => StructValue::from_str(s, st)?.into(),
             DataType::Jsonb => JsonbVal::from_str(s)?.into(),
-            DataType::Bytea => str_to_bytea(s)?.into(),
+            DataType::Bytea => {
+                let mut buf = Vec::new();
+                str_to_bytea(s, &mut buf)?;
+                buf.into()
+            }
             DataType::Vector(size) => VectorVal::from_text(s, *size)?.into(),
             DataType::Map(_m) => return Err("map from text is not supported".into()),
         })

src/common/src/types/mod.rs

@@ -643,7 +643,9 @@ impl JsonParseOptions {
 
                 match self.bytea_handling {
                     ByteaHandling::Standard => {
-                        str_to_bytea(value_str).map_err(|_| create_error())?.into()
+                        let mut buf = Vec::new();
+                        str_to_bytea(value_str, &mut buf).map_err(|_| create_error())?;
+                        buf.into()
                     }
                     ByteaHandling::Base64 => base64::engine::general_purpose::STANDARD
                         .decode(value_str)

src/connector/src/parser/unified/json.rs

@@ -74,6 +74,7 @@ workspace-hack = { path = "../../workspace-hack" }
 [dev-dependencies]
 criterion = { workspace = true }
 expect-test = "1"
+hex-literal = "1.1.0"
 tokio = { version = "0.2", package = "madsim-tokio", features = [
     "rt",
     "macros",

src/expr/impl/Cargo.toml

@@ -12,8 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use std::fmt::Write;
-
 use risingwave_common::array::ListRef;
 use risingwave_expr::{ExprError, Result, function};
 
@@ -185,6 +183,6 @@ fn array_length_of_dim(array: ListRef<'_>, d: i32) -> Result<Option<i32>> {
 /// select array_dims(array[array[]::int[]]); -- would be `[1:1][1:0]` after multidimensional support
 /// ```
 #[function("array_dims(anyarray) -> varchar")]
-fn array_dims(array: ListRef<'_>, writer: &mut impl Write) {
+fn array_dims(array: ListRef<'_>, writer: &mut impl std::fmt::Write) {
     write!(writer, "[1:{}]", array.len()).unwrap();
 }