mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-11-14 06:46:56 +00:00
fix: return errors instead of panics for unsupported arrow field types
Struct-typed query results with arrow field types greptimedb cannot represent (e.g. Decimal256) used to panic during schema conversion and row extraction, dropping the client connection. They now surface as query errors: - ConcreteDataType::try_from builds struct types fallibly via the new StructType::try_from_arrow_fields - Value::try_from(ScalarValue::Struct) uses the same fallible path - new try_value_from_array converts an arrow element to Value with error propagation, used by the postgres struct encoding Signed-off-by: Ning Sun <sunning@greptime.com>
This commit is contained in:
@@ -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(_)
|
||||
|
||||
@@ -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<Self> {
|
||||
let struct_fields = fields
|
||||
.iter()
|
||||
.map(|field| {
|
||||
Ok(StructField::new(
|
||||
field.name().clone(),
|
||||
ConcreteDataType::try_from(field.data_type())?,
|
||||
field.is_nullable(),
|
||||
))
|
||||
})
|
||||
.collect::<Result<Vec<_>>>()?;
|
||||
Ok(StructType {
|
||||
fields: Arc::new(struct_fields),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn fields(&self) -> Arc<Vec<StructField>> {
|
||||
self.fields.clone()
|
||||
}
|
||||
|
||||
@@ -1228,7 +1228,9 @@ impl TryFrom<ScalarValue> 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<ScalarValue> 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<Value> {
|
||||
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<ValueRef<'_>> 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));
|
||||
|
||||
@@ -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<S: Encoder>(
|
||||
|
||||
/// 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<S: Encoder>(
|
||||
query_ctx: &QueryContextRef,
|
||||
column: &ArrayRef,
|
||||
@@ -267,8 +271,8 @@ fn encode_structured_json<S: Encoder>(
|
||||
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<S, B>
|
||||
@@ -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 {
|
||||
|
||||
@@ -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();
|
||||
|
||||
|
||||
Reference in New Issue
Block a user