Add Decimal128 support. (returnString#297)

Author: mshauneu

.gitignore

@@ -1,2 +1,3 @@
 /target
 Cargo.lock
+.idea/

convergence-arrow/Cargo.toml

@@ -13,6 +13,5 @@ async-trait = "0.1"
 datafusion = "43"
 convergence = { path = "../convergence", version = "0.16.0" }
 chrono = "0.4"
-
-[dev-dependencies]
 tokio-postgres = { version =  "0.7", features = [ "with-chrono-0_4" ] }
+rust_decimal = { version = "1.37.1", features = ["default", "db-postgres"] }

convergence-arrow/src/table.rs

@@ -3,9 +3,10 @@
 use convergence::protocol::{DataTypeOid, ErrorResponse, FieldDescription, SqlState};
 use convergence::protocol_ext::DataRowBatch;
 use datafusion::arrow::array::{
-	BooleanArray, Date32Array, Date64Array, Float16Array, Float32Array, Float64Array, Int16Array, Int32Array,
-	Int64Array, Int8Array, StringArray, StringViewArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray,
-	TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array, UInt8Array,
+	BooleanArray, Date32Array, Date64Array, Decimal128Array, Float16Array, Float32Array, Float64Array, Int16Array,
+	Int32Array, Int64Array, Int8Array, StringArray, StringViewArray, TimestampMicrosecondArray,
+	TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, UInt16Array, UInt32Array, UInt64Array,
+	UInt8Array,
 };
 use datafusion::arrow::datatypes::{DataType, Schema, TimeUnit};
 use datafusion::arrow::record_batch::RecordBatch;
@@ -47,6 +48,7 @@ pub fn record_batch_to_rows(arrow_batch: &RecordBatch, pg_batch: &mut DataRowBat
 					DataType::Float16 => row.write_float4(array_val!(Float16Array, col, row_idx).to_f32()),
 					DataType::Float32 => row.write_float4(array_val!(Float32Array, col, row_idx)),
 					DataType::Float64 => row.write_float8(array_val!(Float64Array, col, row_idx)),
+					DataType::Decimal128(p, s) => row.write_numeric_16(array_val!(Decimal128Array, col, row_idx), p, s),
 					DataType::Utf8 => row.write_string(array_val!(StringArray, col, row_idx)),
 					DataType::Utf8View => row.write_string(array_val!(StringViewArray, col, row_idx)),
 					DataType::Date32 => {
@@ -103,8 +105,8 @@ pub fn data_type_to_oid(ty: &DataType) -> Result<DataTypeOid, ErrorResponse> {
 		DataType::UInt64 => DataTypeOid::Int8,
 		DataType::Float16 | DataType::Float32 => DataTypeOid::Float4,
 		DataType::Float64 => DataTypeOid::Float8,
-		DataType::Utf8 => DataTypeOid::Text,
-		DataType::Utf8View => DataTypeOid::Text,
+		DataType::Decimal128(_, _) => DataTypeOid::Numeric,
+		DataType::Utf8 | DataType::Utf8View => DataTypeOid::Text,
 		DataType::Date32 | DataType::Date64 => DataTypeOid::Date,
 		DataType::Timestamp(_, None) => DataTypeOid::Timestamp,
 		other => {

convergence-arrow/tests/test_arrow.rs

@@ -6,10 +6,11 @@ use convergence::protocol_ext::DataRowBatch;
 use convergence::server::{self, BindOptions};
 use convergence::sqlparser::ast::Statement;
 use convergence_arrow::table::{record_batch_to_rows, schema_to_field_desc};
-use datafusion::arrow::array::{ArrayRef, Date32Array, Float32Array, Int32Array, StringArray, TimestampSecondArray};
+use datafusion::arrow::array::{ArrayRef, Date32Array, Decimal128Array, Float32Array, Int32Array, StringArray, StringViewArray, TimestampSecondArray};
 use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit};
 use datafusion::arrow::record_batch::RecordBatch;
 use std::sync::Arc;
+use rust_decimal::Decimal;
 use tokio_postgres::{connect, NoTls};
 
 struct ArrowPortal {
@@ -31,20 +32,22 @@ impl ArrowEngine {
 	fn new() -> Self {
 		let int_col = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
 		let float_col = Arc::new(Float32Array::from(vec![1.5, 2.5, 3.5])) as ArrayRef;
+		let decimal_col = Arc::new(Decimal128Array::from(vec![11, 22, 33]).with_precision_and_scale(2, 0).unwrap()) as ArrayRef;
 		let string_col = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef;
 		let ts_col = Arc::new(TimestampSecondArray::from(vec![1577836800, 1580515200, 1583020800])) as ArrayRef;
 		let date_col = Arc::new(Date32Array::from(vec![0, 1, 2])) as ArrayRef;
 
 		let schema = Schema::new(vec![
 			Field::new("int_col", DataType::Int32, true),
 			Field::new("float_col", DataType::Float32, true),
+			Field::new("decimal_col", DataType::Decimal128(2, 0), true),
 			Field::new("string_col", DataType::Utf8, true),
 			Field::new("ts_col", DataType::Timestamp(TimeUnit::Second, None), true),
 			Field::new("date_col", DataType::Date32, true),
 		]);
 
 		Self {
-			batch: RecordBatch::try_new(Arc::new(schema), vec![int_col, float_col, string_col, ts_col, date_col])
+			batch: RecordBatch::try_new(Arc::new(schema), vec![int_col, float_col, decimal_col, string_col, string_view_col, ts_col, date_col])
 				.expect("failed to create batch"),
 		}
 	}
@@ -89,8 +92,8 @@ async fn basic_data_types() {
 	let rows = client.query("select 1", &[]).await.unwrap();
 	let get_row = |idx: usize| {
 		let row = &rows[idx];
-		let cols: (i32, f32, &str, NaiveDateTime, NaiveDate) =
-			(row.get(0), row.get(1), row.get(2), row.get(3), row.get(4));
+		let cols: (i32, f32, Decimal, &str, &str, NaiveDateTime, NaiveDate) =
+			(row.get(0), row.get(1), row.get(2), row.get(3), row.get(4), row.get(5), row.get(6));
 		cols
 	};
 
@@ -99,6 +102,7 @@ async fn basic_data_types() {
 		(
 			1,
 			1.5,
+			Decimal::from(11),
 			"a",
 			NaiveDate::from_ymd_opt(2020, 1, 1)
 				.unwrap()
@@ -112,6 +116,7 @@ async fn basic_data_types() {
 		(
 			2,
 			2.5,
+			Decimal::from(22),
 			"b",
 			NaiveDate::from_ymd_opt(2020, 2, 1)
 				.unwrap()
@@ -125,6 +130,7 @@ async fn basic_data_types() {
 		(
 			3,
 			3.5,
+			Decimal::from(33),
 			"c",
 			NaiveDate::from_ymd_opt(2020, 3, 1)
 				.unwrap()

convergence/Cargo.toml

@@ -16,6 +16,5 @@ futures = "0.3"
 sqlparser = "0.46"
 async-trait = "0.1"
 chrono = "0.4"
-
-[dev-dependencies]
+rust_decimal = { version = "1.37.1", features = ["default", "db-postgres"] }
 tokio-postgres = "0.7"

convergence/src/protocol.rs

@@ -75,6 +75,8 @@ data_types! {
 	Float4 = 700, 4
 	Float8 = 701, 8
 
+	Numeric = 1700, -1
+
 	Date = 1082, 4
 	Timestamp = 1114, 8
 

convergence/src/protocol_ext.rs

@@ -3,6 +3,8 @@
 use crate::protocol::{ConnectionCodec, FormatCode, ProtocolError, RowDescription};
 use bytes::{BufMut, BytesMut};
 use chrono::{NaiveDate, NaiveDateTime};
+use rust_decimal::Decimal;
+use tokio_postgres::types::{ToSql, Type};
 use tokio_util::codec::Encoder;
 
 /// Supports batched rows for e.g. returning portal result sets.
@@ -131,6 +133,24 @@ impl<'a> DataRowWriter<'a> {
 		}
 	}
 
+	/// Writes a numeric value for the next column.
+	pub fn write_numeric_16(&mut self, val: i128, _p: &u8, s: &i8) {
+		let decimal = Decimal::from_i128_with_scale(val, *s as u32);
+		match self.parent.format_code {
+			FormatCode::Text => {
+				self.write_string(&decimal.to_string())
+			}
+			FormatCode::Binary => {
+				let numeric_type = Type::from_oid(1700).expect("failed to create numeric type");
+				let mut buf = BytesMut::new();
+				decimal.to_sql(&numeric_type, &mut buf)
+					.expect("failed to write numeric");
+
+				self.write_value(&buf.freeze())
+			}
+		};
+	}
+
 	primitive_write!(write_int2, i16);
 	primitive_write!(write_int4, i32);
 	primitive_write!(write_int8, i64);