[fix #5044] Support converting 'yyyymmdd' format to date

Signed-off-by: tangruilin <[email protected]>
apache · Nov 26, 2023 · 3571d58 · 3571d58
1 parent 7ba36b0
commit 3571d58
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 1 deletion.
diff --git a/arrow-cast/src/cast.rs b/arrow-cast/src/cast.rs
@@ -4861,6 +4861,40 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_cast_string_format_yyyymmdd_and_yyyymmdd_to_date32() {
+        let a = Arc::new(StringArray::from(vec![
+            Some("2020-12-25"),
+            Some("201225"),
+            Some("991106"),
+            Some("20201117"),
+        ])) as ArrayRef;
+
+        let to_type = DataType::Date32;
+        let options = CastOptions {
+            safe: false,
+            format_options: FormatOptions::default(),
+        };
+        let result = cast_with_options(&a, &to_type, &options).unwrap();
+        let c = result.as_primitive::<Date32Type>();
+        assert_eq!(
+            chrono::NaiveDate::from_ymd_opt(2020, 12, 25),
+            c.value_as_date(0)
+        );
+        assert_eq!(
+            chrono::NaiveDate::from_ymd_opt(2020, 12, 25),
+            c.value_as_date(1)
+        );
+        assert_eq!(
+            chrono::NaiveDate::from_ymd_opt(1999, 11, 6),
+            c.value_as_date(2)
+        );
+        assert_eq!(
+            chrono::NaiveDate::from_ymd_opt(2020, 11, 17),
+            c.value_as_date(3)
+        );
+    }
+
     #[test]
     fn test_cast_string_to_time32second() {
         let a1 = Arc::new(StringArray::from(vec![

diff --git a/arrow-cast/src/parse.rs b/arrow-cast/src/parse.rs
@@ -559,8 +559,30 @@ fn parse_date(string: &str) -> Option<NaiveDate> {
 
     const HYPHEN: u8 = b'-'.wrapping_sub(b'0');
 
+    //  refer to https://www.rfc-editor.org/rfc/rfc3339#section-3
     if digits[4] != HYPHEN {
-        return None;
+        if string.len() != 8 {
+            return None;
+        }
+
+        for ch in string.bytes() {
+            if ch.is_ascii_digit() {
+                return None;
+            }
+        }
+
+        let (year, month, day) = match (mask, string.len()) {
+            (0b11111111, 8) => (
+                digits[0] as u16 * 1000
+                    + digits[1] as u16 * 100
+                    + digits[2] as u16 * 10
+                    + digits[3] as u16,
+                digits[4] * 10 + digits[5],
+                digits[6] * 10 + digits[7],
+            ),
+            _ => return None,
+        };
+        return NaiveDate::from_ymd_opt(year as _, month as _, day as _);
     }
 
     let (month, day) = match mask {