feat: add support for decimal and decimal64 types in taosws with corresponding tests

Author: qevolg

taos-ws-py/Cargo.lock

@@ -11,6 +11,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 anyhow = "1"
+bigdecimal = "0.4"
 chrono = "0.4"
 chrono-tz = "0.10.4"
 iana-time-zone = "0.1.63"

taos-ws-py/Cargo.toml

@@ -3,12 +3,15 @@
 use std::str::FromStr;
 
 use ::taos::{sync::*, RawBlock, ResultSet};
+use bigdecimal::BigDecimal;
 use chrono_tz::Tz;
 use pyo3::prelude::*;
-use pyo3::types::{PyDict, PyString, PyTuple};
+use pyo3::types::{PyAny, PyDict, PyString, PyTuple, PyType};
 use pyo3::{create_exception, exceptions::PyException};
 use taos::taos_query;
-use taos::taos_query::common::{SchemalessPrecision, SchemalessProtocol, SmlDataBuilder};
+use taos::taos_query::common::{
+    views::DecimalView, SchemalessPrecision, SchemalessProtocol, SmlDataBuilder,
+};
 use taos::Value::{
     BigInt, Bool, Double, Float, Geometry, Int, Json, NChar, Null, SmallInt, Timestamp, TinyInt,
     UBigInt, UInt, USmallInt, UTinyInt, VarBinary, VarChar,
@@ -618,6 +621,7 @@ enum PyColumnType {
     Json,
     VarBinary,
     Decimal,
+    Decimal64,
     Blob,
     MediumBlob,
     Geometry,
@@ -935,6 +939,70 @@ fn doubles_to_column(values: Vec<Option<f64>>) -> PyColumnView {
     }
 }
 
+fn parse_decimal_values(values: Vec<Option<&PyAny>>) -> PyResult<Vec<Option<BigDecimal>>> {
+    let mut decimals = Vec::with_capacity(values.len());
+
+    for (index, value) in values.into_iter().enumerate() {
+        let decimal = match value {
+            Some(value) => {
+                let decimal_type: &PyType = value
+                    .py()
+                    .import("decimal")?
+                    .getattr("Decimal")?
+                    .downcast()?;
+                if !value.is_instance(decimal_type)? {
+                    let type_name = value.get_type().name()?;
+                    return Err(ProgrammingError::new_err(format!(
+                        "expected decimal.Decimal or None, got {type_name} at index {index}"
+                    )));
+                }
+
+                let decimal_str = value.str()?.to_string();
+                let is_finite = value.call_method0("is_finite")?.extract::<bool>()?;
+                if !is_finite {
+                    return Err(ProgrammingError::new_err(format!(
+                        "decimal value must be finite, got '{decimal_str}' at index {index}"
+                    )));
+                }
+
+                let fixed_decimal = value
+                    .call_method1("__format__", ("f",))?
+                    .extract::<String>()?;
+                let parsed = BigDecimal::from_str(&fixed_decimal).map_err(|_| {
+                    ProgrammingError::new_err(format!(
+                        "failed to parse decimal value '{decimal_str}' at index {index}"
+                    ))
+                })?;
+                Some(parsed)
+            }
+            None => None,
+        };
+        decimals.push(decimal);
+    }
+
+    Ok(decimals)
+}
+
+#[pyfunction]
+fn decimal64_to_column(values: Vec<Option<&PyAny>>) -> PyResult<PyColumnView> {
+    let decimals = parse_decimal_values(values)?;
+    let decimal_view = DecimalView::<i64>::from_decimals(decimals)
+        .map_err(|_| ProgrammingError::new_err("decimal64 scale exceeds maximum 18"))?;
+    Ok(PyColumnView {
+        _inner: ColumnView::Decimal64(decimal_view),
+    })
+}
+
+#[pyfunction]
+fn decimal_to_column(values: Vec<Option<&PyAny>>) -> PyResult<PyColumnView> {
+    let decimals = parse_decimal_values(values)?;
+    let decimal_view = DecimalView::<i128>::from_decimals(decimals)
+        .map_err(|_| ProgrammingError::new_err("decimal scale exceeds maximum 38"))?;
+    Ok(PyColumnView {
+        _inner: ColumnView::Decimal(decimal_view),
+    })
+}
+
 #[pyfunction]
 fn varchar_to_column(values: Vec<Option<String>>) -> PyColumnView {
     PyColumnView {
@@ -1073,6 +1141,8 @@ fn _taosws(py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(unsigned_big_ints_to_column, m)?)?;
     m.add_function(wrap_pyfunction!(floats_to_column, m)?)?;
     m.add_function(wrap_pyfunction!(doubles_to_column, m)?)?;
+    m.add_function(wrap_pyfunction!(decimal64_to_column, m)?)?;
+    m.add_function(wrap_pyfunction!(decimal_to_column, m)?)?;
     m.add_function(wrap_pyfunction!(varchar_to_column, m)?)?;
     m.add_function(wrap_pyfunction!(nchar_to_column, m)?)?;
     m.add_function(wrap_pyfunction!(json_to_column, m)?)?;

taos-ws-py/src/lib.rs

@@ -1,6 +1,8 @@
 import time
-import taosws
+from decimal import Decimal
 
+import pytest
+import taosws
 
 url = "taosws://localhost:6041/"
 
@@ -96,3 +98,80 @@ def test_stmt2_stable():
     assert rows == 4
     stmt2_query(conn, "select * from stb1 where a > ?")
     after_test(db_name)
+
+
+def test_decimal_column_rejects_non_decimal_type():
+    with pytest.raises(taosws.ProgrammingError, match=r"expected decimal\.Decimal or None, got int at index 0"):
+        taosws.decimal64_to_column([1])
+
+
+def test_decimal_column_rejects_non_finite_value():
+    with pytest.raises(taosws.ProgrammingError, match=r"decimal value must be finite, got 'NaN' at index 0"):
+        taosws.decimal_to_column([Decimal("NaN")])
+
+
+def test_decimal64_column_rejects_scale_overflow():
+    with pytest.raises(taosws.ProgrammingError, match=r"decimal64 scale exceeds maximum 18"):
+        taosws.decimal64_to_column([Decimal("0.1234567890123456789")])
+
+
+def test_decimal_column_rejects_scale_overflow():
+    with pytest.raises(taosws.ProgrammingError, match=r"decimal scale exceeds maximum 38"):
+        taosws.decimal_to_column([Decimal("0.123456789012345678901234567890123456789")])
+
+
+def test_stmt2_decimal():
+    db_name = "test_1774702104"
+    conn = taosws.connect()
+    try:
+        conn.execute(f"drop database if exists {db_name}")
+        conn.execute(f"create database {db_name}")
+        conn.execute(f"use {db_name}")
+        conn.execute("create table t_decimal (ts timestamp, d64 decimal(10,2), d128 decimal(20,10))")
+
+        stmt2 = conn.stmt2_statement()
+        stmt2.prepare("insert into t_decimal values (?, ?, ?)")
+        param = taosws.stmt2_bind_param_view(
+            table_name="",
+            tags=None,
+            columns=[
+                taosws.millis_timestamps_to_column([1726803356466, 1726803356467, 1726803356468]),
+                taosws.decimal64_to_column([Decimal("99.9876"), Decimal("1.0234"), None]),
+                taosws.decimal_to_column(
+                    [Decimal("1234567890.1234567890"), Decimal("1.23E+5"), Decimal("0.1234567890123")]
+                ),
+            ],
+        )
+        stmt2.bind([param])
+        rows = stmt2.execute()
+        assert rows == 3
+
+        taosws.decimal64_to_column([])
+        taosws.decimal_to_column([])
+
+        try:
+            stmt2.prepare("select d64, d128 from t_decimal where ts >= ? order by ts")
+            query_param = taosws.stmt2_bind_param_view(
+                table_name="",
+                tags=None,
+                columns=[taosws.millis_timestamps_to_column([1726803356466])],
+            )
+            stmt2.bind([query_param])
+            stmt2.execute()
+            data = [row for row in stmt2.result_set()]
+        except (taosws.QueryError, taosws.OperationalError) as err:
+            if "[0x073A]" in str(err):
+                pytest.skip("current ws environment cannot query DECIMAL rows: query memory exhausted")
+            raise
+
+        assert len(data) == 3
+        assert Decimal(data[0][0]) == Decimal("99.99")
+        assert Decimal(data[0][1]) == Decimal("1234567890.1234567890")
+        assert Decimal(data[1][0]) == Decimal("1.02")
+        assert Decimal(data[1][1]) == Decimal("123000")
+        assert data[2][0] is None
+        assert Decimal(data[2][1]) == Decimal("0.1234567890")
+
+    finally:
+        conn.execute(f"drop database if exists {db_name}")
+        conn.close()