Add support for parsing FixedSizeList type

apache · Jan 15, 2024 · b8b12a1 · b8b12a1
1 parent ff728d6
commit b8b12a1
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 0 deletions.
diff --git a/datafusion/sql/src/expr/arrow_cast.rs b/datafusion/sql/src/expr/arrow_cast.rs
@@ -150,6 +150,7 @@ impl<'a> Parser<'a> {
             Token::Dictionary => self.parse_dictionary(),
             Token::List => self.parse_list(),
             Token::LargeList => self.parse_large_list(),
+            Token::FixedSizeList => self.parse_fixed_size_list(),
             tok => Err(make_error(
                 self.val,
                 &format!("finding next type, got unexpected '{tok}'"),
@@ -177,6 +178,18 @@ impl<'a> Parser<'a> {
         ))))
     }
 
+    /// Parses the FixedSizeList type
+    fn parse_fixed_size_list(&mut self) -> Result<DataType> {
+        self.expect_token(Token::LParen)?;
+        let length = self.parse_i32("FixedSizeList")?;
+        self.expect_token(Token::Comma)?;
+        let data_type = self.parse_next_type()?;
+        self.expect_token(Token::RParen)?;
+        Ok(DataType::FixedSizeList(Arc::new(Field::new(
+            "item", data_type, true,
+        )), length))
+    }
+
     /// Parses the next timeunit
     fn parse_time_unit(&mut self, context: &str) -> Result<TimeUnit> {
         match self.next_token()? {
@@ -508,6 +521,7 @@ impl<'a> Tokenizer<'a> {
 
             "List" => Token::List,
             "LargeList" => Token::LargeList,
+            "FixedSizeList" => Token::FixedSizeList,
 
             "Second" => Token::TimeUnit(TimeUnit::Second),
             "Millisecond" => Token::TimeUnit(TimeUnit::Millisecond),
@@ -598,6 +612,7 @@ enum Token {
     DoubleQuotedString(String),
     List,
     LargeList,
+    FixedSizeList,
 }
 
 impl Display for Token {
@@ -606,6 +621,7 @@ impl Display for Token {
             Token::SimpleType(t) => write!(f, "{t}"),
             Token::List => write!(f, "List"),
             Token::LargeList => write!(f, "LargeList"),
+            Token::FixedSizeList => write!(f, "FixedSizeList"),
             Token::Timestamp => write!(f, "Timestamp"),
             Token::Time32 => write!(f, "Time32"),
             Token::Time64 => write!(f, "Time64"),

diff --git a/datafusion/sqllogictest/test_files/arrow_typeof.slt b/datafusion/sqllogictest/test_files/arrow_typeof.slt
@@ -376,3 +376,26 @@ query T
 select arrow_typeof(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)'));
 ----
 LargeList(Field { name: "item", data_type: Int64, nullable: true, dict_id: 0, dict_is_ordered: false, metadata: {} })
+
+
+## FixedSizeList
+
+
+query ?
+select arrow_cast('1', 'FixedSizeList(1, Int64)');
+----
+[1]
+
+query ?
+select arrow_cast('1', 'FixedSizeList(2, Int64)');
+----
+[1]
+
+query ?
+select arrow_cast(make_array(1, 2, 3), 'FixedSizeList(5, Int64)');
+----
+[1]
+
+# LargeList to FixedSizeList
+query ?
+select arrow_typeof(arrow_cast(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)')), 'FixedSizeList(5, Int64)');