feat(query): Support Vector data type (#18044)

* feat(query): Support Vector data type

* native format support vector

* fix

* remove impl ArgType

Author: b41sh

src/common/native/src/nested.rs

@@ -18,6 +18,8 @@ use databend_common_expression::types::AnyType;
 use databend_common_expression::types::ArrayColumn;
 use databend_common_expression::types::Bitmap;
 use databend_common_expression::types::Buffer;
+use databend_common_expression::types::NumberColumn;
+use databend_common_expression::types::VectorColumn;
 use databend_common_expression::Column;
 use databend_common_expression::TableDataType;
 
@@ -30,6 +32,8 @@ pub enum Nested {
     Primitive(usize, bool, Option<Bitmap>),
     /// a list
     LargeList(ListNested),
+    /// a fixed list
+    FixedList(FixedListNested),
     /// A struct column
     Struct(usize, bool, Option<Bitmap>),
 }
@@ -51,13 +55,38 @@ impl ListNested {
     }
 }
 
+#[derive(Debug, Clone, PartialEq)]
+pub struct FixedListNested {
+    pub is_nullable: bool,
+    pub dimension: usize,
+    pub length: usize,
+    pub validity: Option<Bitmap>,
+}
+
+impl FixedListNested {
+    pub fn new(
+        dimension: usize,
+        length: usize,
+        validity: Option<Bitmap>,
+        is_nullable: bool,
+    ) -> Self {
+        Self {
+            is_nullable,
+            dimension,
+            length,
+            validity,
+        }
+    }
+}
+
 pub type NestedState = Vec<Nested>;
 
 impl Nested {
     pub fn length(&self) -> usize {
         match self {
             Nested::Primitive(len, _, _) => *len,
             Nested::LargeList(l) => l.offsets.len(),
+            Nested::FixedList(l) => l.length,
             Nested::Struct(len, _, _) => *len,
         }
     }
@@ -66,37 +95,34 @@ impl Nested {
         match self {
             Nested::Primitive(_, b, _) => *b,
             Nested::LargeList(l) => l.is_nullable,
+            Nested::FixedList(l) => l.is_nullable,
             Nested::Struct(_, b, _) => *b,
         }
     }
 
-    pub fn inner(&self) -> (Buffer<u64>, &Option<Bitmap>) {
+    pub fn offsets(&self) -> Option<Buffer<u64>> {
         match self {
-            Nested::Primitive(_, _, v) => (Buffer::new(), v),
             Nested::LargeList(l) => {
                 let start = *l.offsets.first().unwrap();
                 let buffer = if start == 0 {
                     l.offsets.clone()
                 } else {
                     l.offsets.iter().map(|x| *x - start).collect()
                 };
-                (buffer, &l.validity)
+                Some(buffer)
             }
-            Nested::Struct(_, _, v) => (Buffer::new(), v),
+            _ => None,
         }
     }
 
     pub fn validity(&self) -> &Option<Bitmap> {
         match self {
             Nested::Primitive(_, _, v) => v,
             Nested::LargeList(l) => &l.validity,
+            Nested::FixedList(l) => &l.validity,
             Nested::Struct(_, _, v) => v,
         }
     }
-
-    pub fn is_list(&self) -> bool {
-        matches!(self, Nested::LargeList(_))
-    }
 }
 
 /// Constructs the necessary `Vec<Vec<Nested>>` to write the rep and def levels of `column` to parquet
@@ -110,7 +136,10 @@ pub fn to_nested(column: &Column) -> Result<Vec<Vec<Nested>>> {
 pub fn is_nested_type(t: &TableDataType) -> bool {
     matches!(
         t,
-        TableDataType::Tuple { .. } | TableDataType::Array(_) | TableDataType::Map(_)
+        TableDataType::Tuple { .. }
+            | TableDataType::Array(_)
+            | TableDataType::Map(_)
+            | TableDataType::Vector(_)
     )
 }
 
@@ -133,6 +162,13 @@ pub fn slice_nest_column(
                 current_length = r as usize;
                 current_offset = *l_nested.offsets.first().unwrap() as usize;
             }
+            Nested::FixedList(l_nested) => {
+                if let Some(validity) = l_nested.validity.as_mut() {
+                    validity.slice(current_offset, current_length)
+                };
+                current_offset *= l_nested.dimension;
+                current_length *= l_nested.dimension;
+            }
             Nested::Struct(length, _, validity) => {
                 *length = current_length;
                 if let Some(validity) = validity.as_mut() {
@@ -176,6 +212,15 @@ fn to_nested_recursive(
             }));
             to_nested_recursive(&inner.underlying_column(), nested, parents)?;
         }
+        Column::Vector(inner) => {
+            parents.push(Nested::FixedList(FixedListNested {
+                is_nullable: nullable,
+                dimension: inner.dimension(),
+                length: inner.len(),
+                validity,
+            }));
+            to_nested_recursive(&inner.underlying_column(), nested, parents)?;
+        }
         Column::Map(inner) => {
             parents.push(Nested::LargeList(ListNested {
                 is_nullable: nullable,
@@ -208,6 +253,9 @@ fn to_leaves_recursive(column: &Column, leaves: &mut Vec<Column>) {
         Column::Array(col) => {
             to_leaves_recursive(&col.underlying_column(), leaves);
         }
+        Column::Vector(col) => {
+            to_leaves_recursive(&col.underlying_column(), leaves);
+        }
         Column::Map(col) => {
             to_leaves_recursive(&col.underlying_column(), leaves);
         }
@@ -218,14 +266,15 @@ fn to_leaves_recursive(column: &Column, leaves: &mut Vec<Column>) {
     }
 }
 
-/// The initial info of nested data types.
 /// The initial info of nested data types.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 pub enum InitNested {
     /// Primitive data types
     Primitive(bool),
     /// List data types
     List(bool),
+    /// Fixed List data types
+    FixedList(bool),
     /// Struct data types
     Struct(bool),
 }
@@ -235,14 +284,16 @@ impl InitNested {
         match self {
             InitNested::Primitive(b) => *b,
             InitNested::List(b) => *b,
+            InitNested::FixedList(b) => *b,
             InitNested::Struct(b) => *b,
         }
     }
 }
 
 pub fn create_list(data_type: TableDataType, nested: &mut NestedState, values: Column) -> Column {
     let n = nested.pop().unwrap();
-    let (offsets, validity) = n.inner();
+    let offsets = n.offsets().unwrap();
+    let validity = n.validity();
     let col = Column::Array(Box::new(ArrayColumn::<AnyType>::new(values, offsets)));
 
     if data_type.is_nullable() {
@@ -252,10 +303,36 @@ pub fn create_list(data_type: TableDataType, nested: &mut NestedState, values: C
     }
 }
 
+pub fn create_fixed_list(
+    data_type: TableDataType,
+    dimension: usize,
+    nested: &mut NestedState,
+    values: Column,
+) -> Column {
+    let n = nested.pop().unwrap();
+    let validity = n.validity();
+    let col = match values {
+        Column::Number(NumberColumn::Int8(vals)) => {
+            Column::Vector(VectorColumn::Int8((vals.clone(), dimension)))
+        }
+        Column::Number(NumberColumn::Float32(vals)) => {
+            Column::Vector(VectorColumn::Float32((vals.clone(), dimension)))
+        }
+        _ => unreachable!(),
+    };
+
+    if data_type.is_nullable() {
+        col.wrap_nullable(validity.clone())
+    } else {
+        col
+    }
+}
+
 /// Creates a new [`Mapcolumn`].
 pub fn create_map(data_type: TableDataType, nested: &mut NestedState, values: Column) -> Column {
     let n = nested.pop().unwrap();
-    let (offsets, validity) = n.inner();
+    let offsets = n.offsets().unwrap();
+    let validity = n.validity();
     let col = Column::Map(Box::new(ArrayColumn::<AnyType>::new(values, offsets)));
     if data_type.is_nullable() {
         col.wrap_nullable(validity.clone())
@@ -271,7 +348,7 @@ pub fn create_struct(
 ) -> (NestedState, Column) {
     let mut nest = nested.pop().unwrap();
     let n = nest.pop().unwrap();
-    let (_, validity) = n.inner();
+    let validity = n.validity();
 
     let col = Column::Tuple(values);
     if is_nullable {

src/common/native/src/read/array/fixed_list.rs

@@ -0,0 +1,72 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use databend_common_expression::Column;
+use databend_common_expression::TableDataType;
+
+use crate::error::Result;
+use crate::nested::create_fixed_list;
+use crate::nested::NestedState;
+use crate::read::deserialize::DynIter;
+
+/// An iterator adapter over [`DynIter`] assumed to be encoded as List columns
+pub struct FixedListIterator<'a> {
+    iter: DynIter<'a, Result<(NestedState, Column)>>,
+    data_type: TableDataType,
+    dimension: usize,
+}
+
+impl<'a> FixedListIterator<'a> {
+    /// Creates a new [`FixedListIterator`] with `iter` and `data_type`.
+    pub fn new(
+        iter: DynIter<'a, Result<(NestedState, Column)>>,
+        data_type: TableDataType,
+        dimension: usize,
+    ) -> Self {
+        Self {
+            iter,
+            data_type,
+            dimension,
+        }
+    }
+}
+
+impl FixedListIterator<'_> {
+    fn deserialize(
+        &mut self,
+        value: Option<Result<(NestedState, Column)>>,
+    ) -> Option<Result<(NestedState, Column)>> {
+        let (mut nested, values) = match value {
+            Some(Ok((nested, values))) => (nested, values),
+            Some(Err(err)) => return Some(Err(err)),
+            None => return None,
+        };
+        let array = create_fixed_list(self.data_type.clone(), self.dimension, &mut nested, values);
+        Some(Ok((nested, array)))
+    }
+}
+
+impl Iterator for FixedListIterator<'_> {
+    type Item = Result<(NestedState, Column)>;
+
+    fn nth(&mut self, n: usize) -> Option<Self::Item> {
+        let value = self.iter.nth(n);
+        self.deserialize(value)
+    }
+
+    fn next(&mut self) -> Option<Self::Item> {
+        let value = self.iter.next();
+        self.deserialize(value)
+    }
+}

src/common/native/src/read/array/mod.rs

@@ -31,6 +31,8 @@ mod struct_;
 pub use struct_::*;
 mod list;
 pub use list::*;
+mod fixed_list;
+pub use fixed_list::*;
 mod interval;
 mod map;
 pub use interval::*;

src/common/native/src/read/batch_read.rs

@@ -22,6 +22,7 @@ use databend_common_expression::TableDataType;
 use super::array::*;
 use super::NativeReadBuf;
 use crate::error::Result;
+use crate::nested::create_fixed_list;
 use crate::nested::create_list;
 use crate::nested::create_map;
 use crate::nested::create_struct;
@@ -153,6 +154,18 @@ pub fn read_nested_column<R: NativeReadBuf>(
             }
             columns
         }
+        Vector(vector_ty) => {
+            init.push(InitNested::FixedList(is_nullable));
+            let dimension = vector_ty.dimension() as usize;
+            let inner_ty = vector_ty.inner_data_type();
+            let results = read_nested_column(readers, inner_ty, init, page_metas)?;
+            let mut columns = Vec::with_capacity(results.len());
+            for (mut nested, values) in results {
+                let array = create_fixed_list(data_type.clone(), dimension, &mut nested, values);
+                columns.push((nested, array));
+            }
+            columns
+        }
         Map(inner) => {
             init.push(InitNested::List(is_nullable));
             let results = read_nested_column(readers, inner.as_ref().clone(), init, page_metas)?;

src/common/native/src/read/deserialize.rs

@@ -199,6 +199,13 @@ where
             let iter = deserialize_nested(readers, inner.as_ref().clone(), init)?;
             DynIter::new(ListIterator::new(iter, data_type.clone()))
         }
+        TableDataType::Vector(vector_ty) => {
+            init.push(InitNested::FixedList(is_nullable));
+            let dimension = vector_ty.dimension() as usize;
+            let inner_ty = vector_ty.inner_data_type();
+            let iter = deserialize_nested(readers, inner_ty, init)?;
+            DynIter::new(FixedListIterator::new(iter, data_type.clone(), dimension))
+        }
         TableDataType::Map(inner) => {
             init.push(InitNested::List(is_nullable));
             let iter = deserialize_nested(readers, inner.as_ref().clone(), init)?;

src/common/native/src/read/read_basic.rs

@@ -20,6 +20,7 @@ use databend_common_expression::types::Bitmap;
 use super::NativeReadBuf;
 use crate::compression::Compression;
 use crate::error::Result;
+use crate::nested::FixedListNested;
 use crate::nested::InitNested;
 use crate::nested::ListNested;
 use crate::nested::Nested;
@@ -80,6 +81,17 @@ pub fn read_nested<R: NativeReadBuf>(
                         n.is_nullable(),
                     )))
                 }
+                InitNested::FixedList(_) => {
+                    let mut buf = vec![0u8; 4];
+                    let length = read_u32(reader, &mut buf)?;
+                    let dimension = read_u32(reader, &mut buf)?;
+                    results.push(Nested::FixedList(FixedListNested::new(
+                        dimension as usize,
+                        length as usize,
+                        bitmap,
+                        n.is_nullable(),
+                    )))
+                }
                 InitNested::Struct(_) => {
                     results.push(Nested::Struct(leaf_length, n.is_nullable(), bitmap))
                 }