diff --git a/src/datatypes/src/data_type.rs b/src/datatypes/src/data_type.rs index f2ed05534d6..8ad388c2ff7 100644 --- a/src/datatypes/src/data_type.rs +++ b/src/datatypes/src/data_type.rs @@ -492,7 +492,9 @@ impl TryFrom<&ArrowDataType> for ConcreteDataType { ArrowDataType::Decimal128(precision, scale) => { ConcreteDataType::decimal128_datatype(*precision, *scale) } - ArrowDataType::Struct(fields) => ConcreteDataType::Struct(StructType::from(fields)), + ArrowDataType::Struct(fields) => { + ConcreteDataType::Struct(StructType::try_from_arrow_fields(fields)?) + } ArrowDataType::Float16 | ArrowDataType::Date64 | ArrowDataType::FixedSizeBinary(_) diff --git a/src/datatypes/src/types/struct_type.rs b/src/datatypes/src/types/struct_type.rs index ffd1d1c0b25..fdbfd4886b9 100644 --- a/src/datatypes/src/types/struct_type.rs +++ b/src/datatypes/src/types/struct_type.rs @@ -19,6 +19,7 @@ use arrow::datatypes::{DataType as ArrowDataType, Field}; use arrow_schema::Fields; use serde::{Deserialize, Serialize}; +use crate::error::Result; use crate::prelude::{ConcreteDataType, DataType, LogicalTypeId}; use crate::value::Value; use crate::vectors::StructVectorBuilder; @@ -100,6 +101,24 @@ impl StructType { } } + /// Like the `From<&Fields>` impl, but returns an error for arrow field + /// types greptimedb cannot represent instead of panicking. + pub fn try_from_arrow_fields(fields: &Fields) -> Result { + let struct_fields = fields + .iter() + .map(|field| { + Ok(StructField::new( + field.name().clone(), + ConcreteDataType::try_from(field.data_type())?, + field.is_nullable(), + )) + }) + .collect::>>()?; + Ok(StructType { + fields: Arc::new(struct_fields), + }) + } + pub fn fields(&self) -> Arc> { self.fields.clone() } diff --git a/src/datatypes/src/value.rs b/src/datatypes/src/value.rs index 2a72516d5db..da1df882aba 100644 --- a/src/datatypes/src/value.rs +++ b/src/datatypes/src/value.rs @@ -1228,7 +1228,9 @@ impl TryFrom for Value { if struct_array.is_empty() || struct_array.is_null(0) { return Ok(Value::Null); } - let struct_type = StructType::from(struct_array.fields()); + // Build the struct type fallibly: an unrepresentable arrow + // field type must surface as an error, not a panic. + let struct_type = StructType::try_from_arrow_fields(struct_array.fields())?; let items = struct_array .columns() .iter() @@ -1264,6 +1266,21 @@ impl TryFrom for Value { } } +/// Converts the value at `index` of an arrow array into a [`Value`]. +/// +/// Unlike the unchecked conversion in [`crate::vectors::StructVector`]'s `get`, +/// this propagates conversion errors (for example arrow types greptimedb +/// cannot represent, such as `Decimal256`) instead of panicking, so callers +/// can surface a proper error to the user. +pub fn try_value_from_array(array: &dyn Array, index: usize) -> Result { + if array.is_null(index) { + return Ok(Value::Null); + } + let scalar = + ScalarValue::try_from_array(array, index).context(ConvertArrowArrayToScalarsSnafu)?; + Value::try_from(scalar) +} + impl From> for Value { fn from(value: ValueRef<'_>) -> Self { match value { @@ -2187,6 +2204,50 @@ pub(crate) mod tests { assert_eq!(Value::Null, null_struct.try_into().unwrap()); } + #[test] + fn test_try_value_from_array() { + use arrow_array::Int32Array; + use datafusion_common::arrow::datatypes::i256; + + let array = Int32Array::from(vec![Some(1), None]); + assert_eq!(Value::Int32(1), try_value_from_array(&array, 0).unwrap()); + assert_eq!(Value::Null, try_value_from_array(&array, 1).unwrap()); + + // a struct with representable fields converts + let supported = ScalarStructBuilder::new() + .with_name_and_scalar("i", ScalarValue::Int32(Some(7))) + .build() + .unwrap(); + let ScalarValue::Struct(array) = supported else { + unreachable!(); + }; + assert_eq!( + Value::Struct(StructValue::new( + vec![Value::Int32(7)], + StructType::new(std::sync::Arc::new(vec![StructField::new( + "i", + ConcreteDataType::int32_datatype(), + false, + )])) + )), + try_value_from_array(array.as_ref(), 0).unwrap() + ); + + // an unsupported arrow field type errors instead of panicking + let unsupported = ScalarStructBuilder::new() + .with_name_and_scalar( + "d", + ScalarValue::Decimal256(Some(i256::from_i128(1)), 38, 10), + ) + .build() + .unwrap(); + let ScalarValue::Struct(array) = unsupported else { + unreachable!(); + }; + let error = try_value_from_array(array.as_ref(), 0).unwrap_err(); + assert!(error.to_string().contains("Unsupported arrow data type")); + } + #[test] fn test_value_from_inner() { assert_eq!(Value::Boolean(true), Value::from(true)); diff --git a/src/servers/src/postgres/types.rs b/src/servers/src/postgres/types.rs index 669f77b1892..c83debd39a1 100644 --- a/src/servers/src/postgres/types.rs +++ b/src/servers/src/postgres/types.rs @@ -37,7 +37,7 @@ use datatypes::json::JsonSettings; use datatypes::prelude::{ConcreteDataType, DataType as _, Value}; use datatypes::schema::{Schema, SchemaRef}; use datatypes::types::{Decimal128Type, IntervalType, TimestampType, jsonb_to_string}; -use datatypes::vectors::Helper; +use datatypes::value::try_value_from_array; use futures::Stream; use pg_interval::Interval as PgInterval; use pgwire::api::Type; @@ -260,6 +260,10 @@ fn encode_struct( /// Encodes one row of a structured column (`Struct`, or `List` of `Struct`) as /// a PostgreSQL `json` value. +/// +/// Returns an error for arrow field types greptimedb cannot represent +/// (e.g. `Decimal256`) instead of panicking, so the client receives a proper +/// error response. fn encode_structured_json( query_ctx: &QueryContextRef, column: &ArrayRef, @@ -267,8 +271,8 @@ fn encode_structured_json( encoder: &mut S, pg_field: &FieldInfo, ) -> PgWireResult<()> { - let vector = Helper::try_into_vector(column.clone()).map_err(convert_err)?; - encode_struct(query_ctx, vector.get(i), encoder, pg_field) + let value = try_value_from_array(column.as_ref(), i).map_err(convert_err)?; + encode_struct(query_ctx, value, encoder, pg_field) } pub(crate) struct RecordBatchRowStream @@ -1946,6 +1950,7 @@ mod test { fn test_encode_list_of_struct_as_json() { use datatypes::types::{StructField, StructType}; use datatypes::value::{ListValue, StructValue}; + use datatypes::vectors::Helper; use pgwire::messages::data::DataRow; // struct type: {a: int32, b: float64} @@ -2052,6 +2057,7 @@ mod test { use common_query::native_histogram::{ encode_native_histogram, native_histogram_column_schema, native_histogram_value_type, }; + use datatypes::vectors::Helper; use pgwire::messages::data::DataRow; fn histogram_value(histogram: &Histogram) -> Value { diff --git a/tests-integration/tests/sql.rs b/tests-integration/tests/sql.rs index 13b51c6cab9..e3b66fc4447 100644 --- a/tests-integration/tests/sql.rs +++ b/tests-integration/tests/sql.rs @@ -1084,6 +1084,25 @@ pub async fn test_postgres_struct_types(store_type: StorageType) { let row = query_one(&client, "SELECT [struct(1), NULL]").await; assert_eq!(row, "[{\"c0\":1},null]"); + // Unsupported arrow field types surface as query errors instead of + // dropping the connection. + let error = client + .simple_query("SELECT struct(arrow_cast('1', 'Decimal256(38, 10)'))") + .await + .unwrap_err(); + let message = match error.as_db_error() { + Some(db_error) => db_error.message().to_string(), + None => error.to_string(), + }; + assert!( + message.contains("Unsupported arrow data type"), + "unexpected error message: {message}" + ); + + // The connection stays usable after the error. + let row = query_one(&client, "SELECT struct(arrow_cast('abc', 'Utf8View'))").await; + assert_eq!(row, "{\"c0\":\"abc\"}"); + drop(client); rx.await.unwrap();