From e26292b558584b4a66614625e5c60ef294cdee34 Mon Sep 17 00:00:00 2001 From: Eric Fu Date: Tue, 19 Nov 2024 16:57:17 +0800 Subject: [PATCH] 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 | 57 +++++++++++++++++++++ e2e_test/extended_mode/README.md | 5 ++ e2e_test/extended_mode/type.slt | 8 +-- src/common/src/types/jsonb.rs | 4 +- src/common/src/types/mod.rs | 2 +- src/common/src/types/num256.rs | 7 +-- src/common/src/types/to_binary.rs | 60 +++++++++++++++++++---- src/tests/e2e_extended_mode/README.md | 10 ++-- 8 files changed, 126 insertions(+), 27 deletions(-) create mode 100644 e2e_test/extended_mode/1dim_list.slt.part create mode 100644 e2e_test/extended_mode/README.md diff --git a/e2e_test/extended_mode/1dim_list.slt.part b/e2e_test/extended_mode/1dim_list.slt.part new file mode 100644 index 0000000000000..1e067bd6bfebd --- /dev/null +++ b/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; \ No newline at end of file diff --git a/e2e_test/extended_mode/README.md b/e2e_test/extended_mode/README.md new file mode 100644 index 0000000000000..2c96d7e0c2326 --- /dev/null +++ b/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' +``` diff --git a/e2e_test/extended_mode/type.slt b/e2e_test/extended_mode/type.slt index b172fcf389abc..c6e5c51b5c048 100644 --- a/e2e_test/extended_mode/type.slt +++ b/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 diff --git a/src/common/src/types/jsonb.rs b/src/common/src/types/jsonb.rs index 9e6fdf8641cb7..4b25741fbe96a 100644 --- a/src/common/src/types/jsonb.rs +++ b/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> { - Ok(Some(self.value_serialize().into())) + ) -> super::to_binary::Result { + Ok(self.value_serialize().into()) } } diff --git a/src/common/src/types/mod.rs b/src/common/src/types/mod.rs index 44be87116643c..ad516eab101c6 100644 --- a/src/common/src/types/mod.rs +++ b/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 { 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 { diff --git a/src/common/src/types/num256.rs b/src/common/src/types/num256.rs index 6c96b3ddbbec8..eccb0a9741ea3 100644 --- a/src/common/src/types/num256.rs +++ b/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> { + fn to_binary_with_type(&self, _ty: &DataType) -> super::to_binary::Result { let mut output = bytes::BytesMut::new(); let buffer = self.to_be_bytes(); output.put_slice(&buffer); - Ok(Some(output.freeze())) + Ok(output.freeze()) } } diff --git a/src/common/src/types/to_binary.rs b/src/common/src/types/to_binary.rs index da7f75f0a2a3f..7c5e88dbc10cf 100644 --- a/src/common/src/types/to_binary.rs +++ b/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 = std::result::Result; /// [`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>; + fn to_binary_with_type(&self, ty: &DataType) -> Result; } 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> { + fn to_binary_with_type(&self, ty: &DataType) -> Result { 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 { + // 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> { + fn to_binary_with_type(&self, ty: &DataType) -> Result { 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" + ) + } } } } diff --git a/src/tests/e2e_extended_mode/README.md b/src/tests/e2e_extended_mode/README.md index 07e67d68e3a0c..53e44e6ac8e44 100644 --- a/src/tests/e2e_extended_mode/README.md +++ b/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