feat(pgwire): support 1-dimension array in extended mode (#19432)

Co-authored-by: Noel Kwan <47273164+kwannoel@users.noreply.github.com>

e2e_test/extended_mode/1dim_list.slt.part

@@ -0,0 +1,57 @@
+# Test binary format of 1-dimension lists (arrays)
+
+query T
+select ARRAY['foo', 'bar', null];
+----
+{foo,bar,NULL}
+
+query T
+select ARRAY[1,2+3,4*5+1];
+----
+{1,5,21}
+
+query T
+select ARRAY[null];
+----
+{NULL}
+
+statement error
+select ARRAY[];
+
+query T
+select ARRAY[]::int[];
+----
+{}
+
+statement ok
+create table t (v1 int);
+
+statement ok
+insert into t values (1), (2), (3);
+
+query T rowsort
+select ARRAY[1, v1*2] from t;
+----
+{1,2}
+{1,4}
+{1,6}
+
+query T
+select min(ARRAY[1, v1*2]) from t;
+----
+{1,2}
+
+query T
+select max(ARRAY[1, v1*2]) from t;
+----
+{1,6}
+
+query T
+select array[false, false] from t;
+----
+{f,f}
+{f,f}
+{f,f}
+
+statement ok
+drop table t;

e2e_test/extended_mode/README.md

@@ -0,0 +1,5 @@
+# How to run
+
+```shell
+sqllogictest -p 4566 -d dev -e postgres-extended './e2e_test/extended_mode/**/*.slt'
+```

e2e_test/extended_mode/type.slt

@@ -3,10 +3,10 @@
 statement ok
 SET RW_IMPLICIT_FLUSH TO true;
 
-# RisingWave can't support list and struct now so we skip them.
-# include ../batch/types/array.slt.part
-# include ../batch/types/struct.slt.part
-# include ../batch/types/list.slt.part
+include 1dim_list.slt.part
+
+# RisingWave can't support struct now so we skip it.
+# include ../batch/types/struct/*.slt.part
 
 # Sqllogitest can't support binary format bytea type so we skip it.
 # include ../batch/types/bytea.slt.part

src/common/src/types/jsonb.rs

@@ -133,8 +133,8 @@ impl crate::types::to_binary::ToBinary for JsonbRef<'_> {
     fn to_binary_with_type(
         &self,
         _ty: &crate::types::DataType,
-    ) -> super::to_binary::Result<Option<bytes::Bytes>> {
-        Ok(Some(self.value_serialize().into()))
+    ) -> super::to_binary::Result<bytes::Bytes> {
+        Ok(self.value_serialize().into())
     }
 }
 

src/common/src/types/mod.rs

@@ -987,7 +987,7 @@ pub fn hash_datum(datum: impl ToDatumRef, state: &mut impl std::hash::Hasher) {
 impl ScalarRefImpl<'_> {
     pub fn binary_format(&self, data_type: &DataType) -> to_binary::Result<Bytes> {
         use self::to_binary::ToBinary;
-        self.to_binary_with_type(data_type).transpose().unwrap()
+        self.to_binary_with_type(data_type)
     }
 
     pub fn text_format(&self, data_type: &DataType) -> String {

src/common/src/types/num256.rs

@@ -165,14 +165,11 @@ macro_rules! impl_common_for_num256 {
         }
 
         impl ToBinary for $scalar_ref<'_> {
-            fn to_binary_with_type(
-                &self,
-                _ty: &DataType,
-            ) -> super::to_binary::Result<Option<Bytes>> {
+            fn to_binary_with_type(&self, _ty: &DataType) -> super::to_binary::Result<Bytes> {
                 let mut output = bytes::BytesMut::new();
                 let buffer = self.to_be_bytes();
                 output.put_slice(&buffer);
-                Ok(Some(output.freeze()))
+                Ok(output.freeze())
             }
         }
 

src/common/src/types/to_binary.rs

@@ -12,13 +12,14 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use bytes::{Bytes, BytesMut};
+use bytes::{BufMut, Bytes, BytesMut};
 use postgres_types::{ToSql, Type};
 
 use super::{
     DataType, Date, Decimal, Interval, ScalarRefImpl, Serial, Time, Timestamp, Timestamptz, F32,
     F64,
 };
+use crate::array::ListRef;
 use crate::error::NotImplemented;
 
 /// Error type for [`ToBinary`] trait.
@@ -38,19 +39,19 @@ pub type Result<T> = std::result::Result<T, ToBinaryError>;
 /// [`postgres_types::ToSql`] has similar functionality, and most of our types implement
 /// that trait and forward `ToBinary` to it directly.
 pub trait ToBinary {
-    fn to_binary_with_type(&self, ty: &DataType) -> Result<Option<Bytes>>;
+    fn to_binary_with_type(&self, ty: &DataType) -> Result<Bytes>;
 }
 macro_rules! implement_using_to_sql {
     ($({ $scalar_type:ty, $data_type:ident, $accessor:expr } ),* $(,)?) => {
         $(
             impl ToBinary for $scalar_type {
-                fn to_binary_with_type(&self, ty: &DataType) -> Result<Option<Bytes>> {
+                fn to_binary_with_type(&self, ty: &DataType) -> Result<Bytes> {
                     match ty {
                         DataType::$data_type => {
                             let mut output = BytesMut::new();
                             #[allow(clippy::redundant_closure_call)]
                             $accessor(self).to_sql(&Type::ANY, &mut output).map_err(ToBinaryError::ToSql)?;
-                            Ok(Some(output.freeze()))
+                            Ok(output.freeze())
                         },
                         _ => unreachable!(),
                     }
@@ -78,8 +79,45 @@ implement_using_to_sql! {
     { Timestamptz, Timestamptz, |x: &Timestamptz| x.to_datetime_utc() }
 }
 
+impl ToBinary for ListRef<'_> {
+    fn to_binary_with_type(&self, ty: &DataType) -> Result<Bytes> {
+        // Reference: Postgres code `src/backend/utils/adt/arrayfuncs.c`
+        // https://github.com/postgres/postgres/blob/c1c09007e219ae68d1f8428a54baf68ccc1f8683/src/backend/utils/adt/arrayfuncs.c#L1548
+        use crate::row::Row;
+        let element_ty = match ty {
+            DataType::List(ty) => ty.as_ref(),
+            _ => unreachable!(),
+        };
+        if matches!(element_ty, DataType::List(_)) {
+            bail_not_implemented!(
+                issue = 7949,
+                "list with 2 or more dimensions is not supported"
+            )
+        }
+        let mut buf = BytesMut::new();
+        buf.put_i32(1); // Number of dimensions (must be 1)
+        buf.put_i32(1); // Has nulls?
+        buf.put_i32(element_ty.to_oid()); // Element type
+        buf.put_i32(self.len() as i32); // Length of 1st dimension
+        buf.put_i32(0); // Offset of 1st dimension
+        for element in self.iter() {
+            match element {
+                None => {
+                    buf.put_i32(-1); // -1 length means a NULL
+                }
+                Some(value) => {
+                    let data = value.to_binary_with_type(element_ty)?;
+                    buf.put_i32(data.len() as i32); // Length of element
+                    buf.put(data);
+                }
+            }
+        }
+        Ok(buf.into())
+    }
+}
+
 impl ToBinary for ScalarRefImpl<'_> {
-    fn to_binary_with_type(&self, ty: &DataType) -> Result<Option<Bytes>> {
+    fn to_binary_with_type(&self, ty: &DataType) -> Result<Bytes> {
         match self {
             ScalarRefImpl::Int16(v) => v.to_binary_with_type(ty),
             ScalarRefImpl::Int32(v) => v.to_binary_with_type(ty),
@@ -98,11 +136,13 @@ impl ToBinary for ScalarRefImpl<'_> {
             ScalarRefImpl::Time(v) => v.to_binary_with_type(ty),
             ScalarRefImpl::Bytea(v) => v.to_binary_with_type(ty),
             ScalarRefImpl::Jsonb(v) => v.to_binary_with_type(ty),
-            ScalarRefImpl::Struct(_) | ScalarRefImpl::List(_) => bail_not_implemented!(
-                issue = 7949,
-                "the pgwire extended-mode encoding for {ty} is unsupported"
-            ),
-            ScalarRefImpl::Map(_) => todo!(),
+            ScalarRefImpl::List(v) => v.to_binary_with_type(ty),
+            ScalarRefImpl::Struct(_) | ScalarRefImpl::Map(_) => {
+                bail_not_implemented!(
+                    issue = 7949,
+                    "the pgwire extended-mode encoding for {ty} is unsupported"
+                )
+            }
         }
     }
 }

src/tests/e2e_extended_mode/README.md

@@ -1,16 +1,16 @@
 This is a program used for e2e test in extended mode.
 
-## What is difference between it and extended_mode/*.slt in e2e_test
+## What is difference between this and `e2e_test/extended_mode`
+
+For e2e test in extended query mode, there are a few things we can't test in sqllogictest
 
-For e2e test in extended query mode, there are two thing we can't test in sqllogitest
 1. bind parameter
 2. max row number
 3. cancel query
-See [detail](https://www.postgresql.org/docs/15/protocol-flow.html#PROTOCOL-FLOW-PIPELINING:~:text=Once%20a%20portal,count%20is%20ignored)
 
-So before sqllogictest supporting these, we test these function in this program.
+See more details [here](https://www.postgresql.org/docs/15/protocol-flow.html#PROTOCOL-FLOW-PIPELINING:~:text=Once%20a%20portal,count%20is%20ignored).
 
-In the future, we may merge it to e2e_text/extended_query
+Before sqllogictest supports these, we test these functions in this program. In the future, we may merge it to `e2e_test/extended_mode`.
 
 # How to run