Merge pull request #6385 from youngsofun/fmt

fix(query): Add NestedCheckpointReader for input format parser
databendlabs · Jul 1, 2022 · bff565e · bff565e
2 parents e3d3f09 + 2674061
commit bff565e
Show file tree

Hide file tree

Showing 28 changed files with 395 additions and 185 deletions.
diff --git a/common/datavalues/src/types/deserializations/array.rs b/common/datavalues/src/types/deserializations/array.rs
@@ -79,7 +79,11 @@ impl TypeDeserializer for ArrayDeserializer {
         }
     }
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        format: &FormatSettings,
+    ) -> Result<()> {
         reader.must_ignore_byte(b'[')?;
         let mut idx = 0;
         loop {
@@ -107,7 +111,7 @@ impl TypeDeserializer for ArrayDeserializer {
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
         reader.must_ignore_byte(b'[')?;
@@ -137,7 +141,7 @@ impl TypeDeserializer for ArrayDeserializer {
 
     fn de_whole_text(&mut self, reader: &[u8], format: &FormatSettings) -> Result<()> {
         let reader = BufferReader::new(reader);
-        let mut reader = CheckpointReader::new(Box::new(reader));
+        let mut reader = NestedCheckpointReader::new(Box::new(reader));
         self.de_text(&mut reader, format)
     }
 

diff --git a/common/datavalues/src/types/deserializations/boolean.rs b/common/datavalues/src/types/deserializations/boolean.rs
@@ -68,7 +68,11 @@ impl TypeDeserializer for BooleanDeserializer {
         Ok(())
     }
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, _format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        _format: &FormatSettings,
+    ) -> Result<()> {
         let v = if BufferReadExt::ignore_insensitive_bytes(reader, b"true")? {
             Ok(true)
         } else if BufferReadExt::ignore_insensitive_bytes(reader, b"false")? {

diff --git a/common/datavalues/src/types/deserializations/date.rs b/common/datavalues/src/types/deserializations/date.rs
@@ -87,7 +87,7 @@ where
 
     fn de_text_quoted<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         _format: &FormatSettings,
     ) -> Result<()> {
         reader.must_ignore_byte(b'\'')?;
@@ -100,7 +100,11 @@ where
         Ok(())
     }
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, _format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        _format: &FormatSettings,
+    ) -> Result<()> {
         let date = reader.read_date_text()?;
         let days = uniform(date);
         check_date(days.as_i32())?;
@@ -110,7 +114,7 @@ where
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         _format: &FormatSettings,
     ) -> Result<()> {
         let maybe_quote = reader.ignore(|f| f == b'\'' || f == b'"')?;
@@ -126,7 +130,7 @@ where
 
     fn de_text_json<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         _format: &FormatSettings,
     ) -> Result<()> {
         reader.must_ignore_byte(b'"')?;

diff --git a/common/datavalues/src/types/deserializations/mod.rs b/common/datavalues/src/types/deserializations/mod.rs
@@ -63,27 +63,31 @@ pub trait TypeDeserializer: Send + Sync {
 
     fn de_whole_text(&mut self, reader: &[u8], format: &FormatSettings) -> Result<()>;
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, format: &FormatSettings) -> Result<()>;
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        format: &FormatSettings,
+    ) -> Result<()>;
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
         self.de_text(reader, format)
     }
 
     fn de_text_json<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
         self.de_text(reader, format)
     }
 
     fn de_text_quoted<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
         self.de_text(reader, format)

diff --git a/common/datavalues/src/types/deserializations/null.rs b/common/datavalues/src/types/deserializations/null.rs
@@ -58,7 +58,11 @@ impl TypeDeserializer for NullDeserializer {
         Ok(())
     }
 
-    fn de_text<R: BufferRead>(&mut self, _reader: &mut R, _format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        _reader: &mut NestedCheckpointReader<R>,
+        _format: &FormatSettings,
+    ) -> Result<()> {
         self.builder.append_default();
         Ok(())
     }

diff --git a/common/datavalues/src/types/deserializations/nullable.rs b/common/datavalues/src/types/deserializations/nullable.rs
@@ -71,60 +71,84 @@ impl TypeDeserializer for NullableDeserializer {
 
     fn de_text_json<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
+        reader.push_checkpoint();
         if reader.ignore_insensitive_bytes(b"null")? {
+            reader.pop_checkpoint();
             self.de_default(format);
             return Ok(());
         }
+        reader.rollback_to_checkpoint()?;
+        reader.pop_checkpoint();
         self.inner.de_text_json(reader, format)?;
         self.bitmap.push(true);
         Ok(())
     }
 
     // TODO: support null text setting
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        format: &FormatSettings,
+    ) -> Result<()> {
+        reader.push_checkpoint();
         if reader.ignore_insensitive_bytes(b"null")? {
+            reader.pop_checkpoint();
             self.de_default(format);
             return Ok(());
+        } else {
+            reader.rollback_to_checkpoint()?;
+            reader.pop_checkpoint();
+            self.inner.de_text(reader, format)?;
+            self.bitmap.push(true);
         }
-        self.inner.de_text(reader, format)?;
-        self.bitmap.push(true);
+
         Ok(())
     }
 
     fn de_text_quoted<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
+        reader.push_checkpoint();
         if reader.ignore_insensitive_bytes(b"null")? {
+            reader.pop_checkpoint();
             self.de_default(format);
-            return Ok(());
+        } else {
+            reader.rollback_to_checkpoint()?;
+            reader.pop_checkpoint();
+            self.inner.de_text_quoted(reader, format)?;
+            self.bitmap.push(true);
         }
-        self.inner.de_text_quoted(reader, format)?;
-        self.bitmap.push(true);
         Ok(())
     }
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
         if reader.eof()? {
             self.de_default(format);
         } else {
-            reader.ignore_insensitive_bytes(b"null")?;
-            let buffer = reader.fill_buf()?;
-            if !buffer.is_empty() && (buffer[0] == b'\r' || buffer[0] == b'\n' || buffer[0] == b',')
-            {
-                self.de_default(format);
-            } else {
-                self.inner.de_text_csv(reader, format)?;
-                self.bitmap.push(true);
+            reader.push_checkpoint();
+            if reader.ignore_insensitive_bytes(b"null")? {
+                let buffer = reader.fill_buf()?;
+                if !buffer.is_empty()
+                    && (buffer[0] == b'\r' || buffer[0] == b'\n' || buffer[0] == b',')
+                {
+                    self.de_default(format);
+                    reader.pop_checkpoint();
+                    return Ok(());
+                }
             }
+            reader.rollback_to_checkpoint()?;
+            reader.pop_checkpoint();
+            self.inner.de_text_csv(reader, format)?;
+            self.bitmap.push(true);
         }
         Ok(())
     }

diff --git a/common/datavalues/src/types/deserializations/number.rs b/common/datavalues/src/types/deserializations/number.rs
@@ -89,7 +89,11 @@ where
         Ok(())
     }
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, _format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        _format: &FormatSettings,
+    ) -> Result<()> {
         let v: T = if !T::FLOATING {
             reader.read_int_text()
         } else {
@@ -101,7 +105,7 @@ where
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         _settings: &FormatSettings,
     ) -> Result<()> {
         let maybe_quote = reader.ignore(|f| f == b'\'' || f == b'"')?;

diff --git a/common/datavalues/src/types/deserializations/string.rs b/common/datavalues/src/types/deserializations/string.rs
@@ -84,7 +84,11 @@ impl TypeDeserializer for StringDeserializer {
         Ok(())
     }
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, _format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        _format: &FormatSettings,
+    ) -> Result<()> {
         self.buffer.clear();
         reader.read_escaped_string_text(&mut self.buffer)?;
         self.builder.append_value(self.buffer.as_slice());
@@ -93,7 +97,7 @@ impl TypeDeserializer for StringDeserializer {
 
     fn de_text_quoted<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         _format: &FormatSettings,
     ) -> Result<()> {
         self.buffer.clear();
@@ -104,7 +108,7 @@ impl TypeDeserializer for StringDeserializer {
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         settings: &FormatSettings,
     ) -> Result<()> {
         let mut read_buffer = reader.fill_buf()?;

diff --git a/common/datavalues/src/types/deserializations/struct_.rs b/common/datavalues/src/types/deserializations/struct_.rs
@@ -81,7 +81,11 @@ impl TypeDeserializer for StructDeserializer {
         }
     }
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        format: &FormatSettings,
+    ) -> Result<()> {
         reader.must_ignore_byte(b'(')?;
         let mut values = Vec::with_capacity(self.inners.len());
         for (idx, inner) in self.inners.iter_mut().enumerate() {
@@ -100,7 +104,7 @@ impl TypeDeserializer for StructDeserializer {
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        _reader: &mut R,
+        _reader: &mut NestedCheckpointReader<R>,
         _format: &FormatSettings,
     ) -> Result<()> {
         Err(ErrorCode::UnImplement("Unimplement error"))

diff --git a/common/datavalues/src/types/deserializations/timestamp.rs b/common/datavalues/src/types/deserializations/timestamp.rs
@@ -70,7 +70,7 @@ impl TypeDeserializer for TimestampDeserializer {
 
     fn de_text_quoted<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
         reader.must_ignore_byte(b'\'')?;
@@ -92,7 +92,11 @@ impl TypeDeserializer for TimestampDeserializer {
         Ok(())
     }
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        format: &FormatSettings,
+    ) -> Result<()> {
         let ts = reader.read_timestamp_text(&format.timezone)?;
         let micros = ts.timestamp_micros();
         check_timestamp(micros)?;
@@ -102,7 +106,7 @@ impl TypeDeserializer for TimestampDeserializer {
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
         let maybe_quote = reader.ignore(|f| f == b'\'' || f == b'"')?;
@@ -118,7 +122,7 @@ impl TypeDeserializer for TimestampDeserializer {
 
     fn de_text_json<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         format: &FormatSettings,
     ) -> Result<()> {
         reader.must_ignore_byte(b'"')?;

diff --git a/common/datavalues/src/types/deserializations/variant.rs b/common/datavalues/src/types/deserializations/variant.rs
@@ -81,7 +81,11 @@ impl TypeDeserializer for VariantDeserializer {
         Ok(())
     }
 
-    fn de_text<R: BufferRead>(&mut self, reader: &mut R, _format: &FormatSettings) -> Result<()> {
+    fn de_text<R: BufferRead>(
+        &mut self,
+        reader: &mut NestedCheckpointReader<R>,
+        _format: &FormatSettings,
+    ) -> Result<()> {
         self.buffer.clear();
         reader.read_escaped_string_text(&mut self.buffer)?;
         let val = serde_json::from_slice(self.buffer.as_slice())?;
@@ -91,7 +95,7 @@ impl TypeDeserializer for VariantDeserializer {
 
     fn de_text_quoted<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         _format: &FormatSettings,
     ) -> Result<()> {
         self.buffer.clear();
@@ -105,7 +109,7 @@ impl TypeDeserializer for VariantDeserializer {
 
     fn de_text_csv<R: BufferRead>(
         &mut self,
-        reader: &mut R,
+        reader: &mut NestedCheckpointReader<R>,
         _format: &FormatSettings,
     ) -> Result<()> {
         self.buffer.clear();