SNOW-1915469 Basic support for DECFLOAT type (#2167)

Author: sfc-gh-tkissinger

DESCRIPTION.md

@@ -46,9 +46,6 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added a feature to verify if the connection is still good enough to send queries over.
   - Added support for base64-encoded DER private key strings in the `private_key` authentication type.
 
-- v3.12.5(TBD)
-  - Added a feature to limit the sizes of IO-bound ThreadPoolExecutors during PUT and GET commands.
-
 - v3.12.4(December 3,2024)
   - Fixed a bug where multipart uploads to Azure would be missing their MD5 hashes.
   - Fixed a bug where OpenTelemetry header injection would sometimes cause Exceptions to be thrown.

setup.py

@@ -101,6 +101,7 @@ def build_extension(self, ext):
                         "CArrowIterator.cpp",
                         "CArrowTableIterator.cpp",
                         "DateConverter.cpp",
+                        "DecFloatConverter.cpp",
                         "DecimalConverter.cpp",
                         "FixedSizeListConverter.cpp",
                         "FloatConverter.cpp",

src/snowflake/connector/arrow_context.py

@@ -159,3 +159,11 @@ def DECIMAL128_to_decimal(self, int128_bytes: bytes, scale: int) -> decimal.Deci
         digits = [int(digit) for digit in str(int128) if digit != "-"]
         sign = int128 < 0
         return decimal.Decimal((sign, digits, -scale))
+
+    def DECFLOAT_to_decimal(self, exponent: int, significand: bytes) -> decimal.Decimal:
+        # significand is two's complement big endian.
+        significand = int.from_bytes(significand, byteorder="big", signed=True)
+        return decimal.Decimal(significand).scaleb(exponent)
+
+    def DECFLOAT_to_numpy_float64(self, exponent: int, significand: bytes) -> float64:
+        return numpy.float64(self.DECFLOAT_to_decimal(exponent, significand))

src/snowflake/connector/converter.py

@@ -203,6 +203,12 @@ def conv(value: str) -> int64:
 
             return conv
 
+    def _DECFLOAT_numpy_to_python(self, ctx: dict[str, Any]) -> Callable:
+        return numpy.float64
+
+    def _DECFLOAT_to_python(self, ctx: dict[str, Any]) -> Callable:
+        return decimal.Decimal
+
     def _REAL_to_python(self, _: dict[str, str | None] | dict[str, str]) -> Callable:
         return float
 

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowChunkIterator.cpp

@@ -12,6 +12,7 @@
 #include "BinaryConverter.hpp"
 #include "BooleanConverter.hpp"
 #include "DateConverter.hpp"
+#include "DecFloatConverter.hpp"
 #include "DecimalConverter.hpp"
 #include "FixedSizeListConverter.hpp"
 #include "FloatConverter.hpp"
@@ -471,6 +472,12 @@ std::shared_ptr<sf::IColumnConverter> getConverterFromSchema(
       break;
     }
 
+    case SnowflakeType::Type::DECFLOAT: {
+      converter = std::make_shared<sf::DecFloatConverter>(*array, schemaView,
+                                                          *context, useNumpy);
+      break;
+    }
+
     default: {
       std::string errorInfo = Logger::formatString(
           "[Snowflake Exception] unknown snowflake data type : %d", st);

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/DecFloatConverter.cpp

@@ -0,0 +1,87 @@
+
+//
+// Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+//
+
+#include "DecFloatConverter.hpp"
+
+#include <cstring>
+#include <memory>
+
+#include "Python/Helpers.hpp"
+
+namespace sf {
+
+Logger* DecFloatConverter::logger =
+    new Logger("snowflake.connector.DecFloatConverter");
+
+const std::string DecFloatConverter::FIELD_NAME_EXPONENT = "exponent";
+const std::string DecFloatConverter::FIELD_NAME_SIGNIFICAND = "significand";
+
+DecFloatConverter::DecFloatConverter(ArrowArrayView& array,
+                                     ArrowSchemaView& schema, PyObject& context,
+                                     bool useNumpy)
+    : m_context(context),
+      m_array(array),
+      m_exponent(nullptr),
+      m_significand(nullptr),
+      m_useNumpy(useNumpy) {
+  if (schema.schema->n_children != 2) {
+    std::string errorInfo = Logger::formatString(
+        "[Snowflake Exception] arrow schema field number does not match, "
+        "expected 2 but got %d instead",
+        schema.schema->n_children);
+    logger->error(__FILE__, __func__, __LINE__, errorInfo.c_str());
+    PyErr_SetString(PyExc_Exception, errorInfo.c_str());
+    return;
+  }
+  for (int i = 0; i < schema.schema->n_children; i += 1) {
+    ArrowSchema* c_schema = schema.schema->children[i];
+    if (std::strcmp(c_schema->name,
+                    DecFloatConverter::FIELD_NAME_EXPONENT.c_str()) == 0) {
+      m_exponent = m_array.children[i];
+    } else if (std::strcmp(c_schema->name,
+                           DecFloatConverter::FIELD_NAME_SIGNIFICAND.c_str()) ==
+               0) {
+      m_significand = m_array.children[i];
+    }
+  }
+  if (!m_exponent || !m_significand) {
+    std::string errorInfo = Logger::formatString(
+        "[Snowflake Exception] arrow schema field names do not match, "
+        "expected %s and %s, but got %s and %s instead",
+        DecFloatConverter::FIELD_NAME_EXPONENT.c_str(),
+        DecFloatConverter::FIELD_NAME_SIGNIFICAND.c_str(),
+        schema.schema->children[0]->name, schema.schema->children[1]->name);
+    logger->error(__FILE__, __func__, __LINE__, errorInfo.c_str());
+    PyErr_SetString(PyExc_Exception, errorInfo.c_str());
+    return;
+  }
+}
+
+PyObject* DecFloatConverter::toPyObject(int64_t rowIndex) const {
+  if (ArrowArrayViewIsNull(&m_array, rowIndex)) {
+    Py_RETURN_NONE;
+  }
+  int64_t exponent = ArrowArrayViewGetIntUnsafe(m_exponent, rowIndex);
+  ArrowStringView stringView =
+      ArrowArrayViewGetStringUnsafe(m_significand, rowIndex);
+  if (stringView.size_bytes > 16) {
+    std::string errorInfo = Logger::formatString(
+        "[Snowflake Exception] only precisions up to 38 supported. "
+        "Please update to a newer version of the connector.");
+    logger->error(__FILE__, __func__, __LINE__, errorInfo.c_str());
+    PyErr_SetString(PyExc_Exception, errorInfo.c_str());
+    return nullptr;
+  }
+  PyObject* significand =
+      PyBytes_FromStringAndSize(stringView.data, stringView.size_bytes);
+
+  PyObject* result = PyObject_CallMethod(
+      &m_context,
+      m_useNumpy ? "DECFLOAT_to_numpy_float64" : "DECFLOAT_to_decimal", "iS",
+      exponent, significand);
+  Py_XDECREF(significand);
+  return result;
+}
+}  // namespace sf

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/DecFloatConverter.hpp

@@ -0,0 +1,39 @@
+
+//
+// Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+//
+
+#ifndef PC_DECFLOATCONVERTER_HPP
+#define PC_DECFLOATCONVERTER_HPP
+
+#include <memory>
+
+#include "IColumnConverter.hpp"
+#include "logging.hpp"
+#include "nanoarrow.h"
+
+namespace sf {
+
+class DecFloatConverter : public IColumnConverter {
+ public:
+  const static std::string FIELD_NAME_EXPONENT;
+  const static std::string FIELD_NAME_SIGNIFICAND;
+
+  explicit DecFloatConverter(ArrowArrayView& array, ArrowSchemaView& schema,
+                             PyObject& context, bool useNumpy);
+
+  PyObject* toPyObject(int64_t rowIndex) const override;
+
+ private:
+  PyObject& m_context;
+  ArrowArrayView& m_array;
+  ArrowArrayView* m_exponent;
+  ArrowArrayView* m_significand;
+  bool m_useNumpy;
+
+  static Logger* logger;
+};
+
+}  // namespace sf
+
+#endif  // PC_DECFLOATCONVERTER_HPP

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/SnowflakeType.cpp

@@ -17,6 +17,7 @@ std::unordered_map<std::string, SnowflakeType::Type>
         {"DOUBLE PRECISION", SnowflakeType::Type::REAL},
         {"DOUBLE", SnowflakeType::Type::REAL},
         {"FIXED", SnowflakeType::Type::FIXED},
+        {"DECFLOAT", SnowflakeType::Type::DECFLOAT},
         {"FLOAT", SnowflakeType::Type::REAL},
         {"MAP", SnowflakeType::Type::MAP},
         {"OBJECT", SnowflakeType::Type::OBJECT},

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/SnowflakeType.hpp

@@ -33,6 +33,7 @@ class SnowflakeType {
     VARIANT = 15,
     VECTOR = 16,
     MAP = 17,
+    DECFLOAT = 18,
   };
 
   static SnowflakeType::Type snowflakeTypeFromString(std::string str) {

test/integ/test_decfloat.py

@@ -0,0 +1,92 @@
+#!/usr/bin/env python
+#
+# Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+#
+
+from __future__ import annotations
+
+import decimal
+from decimal import Decimal
+
+import numpy
+
+import snowflake.connector
+
+
+def test_decfloat_bindings(conn_cnx):
+    # set required decimal precision
+    decimal.getcontext().prec = 38
+    original_style = snowflake.connector.paramstyle
+    snowflake.connector.paramstyle = "qmark"
+    try:
+        with conn_cnx() as cnx:
+            # test decfloat bindings
+            ret = (
+                cnx.cursor()
+                .execute("select ?", [("DECFLOAT", Decimal("-1234e4000"))])
+                .fetchone()
+            )
+            assert isinstance(ret[0], Decimal)
+            assert ret[0] == Decimal("-1234e4000")
+            ret = cnx.cursor().execute("select ?", [("DECFLOAT", -1e3)]).fetchone()
+            assert isinstance(ret[0], Decimal)
+            assert ret[0] == Decimal("-1e3")
+            # test 38 digits
+            ret = (
+                cnx.cursor()
+                .execute(
+                    "select ?",
+                    [("DECFLOAT", Decimal("12345678901234567890123456789012345678"))],
+                )
+                .fetchone()
+            )
+            assert isinstance(ret[0], Decimal)
+            assert ret[0] == Decimal("12345678901234567890123456789012345678")
+            # test w/o explicit type specification
+            ret = cnx.cursor().execute("select ?", [-1e3]).fetchone()
+            assert isinstance(ret[0], float)
+            ret = cnx.cursor().execute("select ?", [Decimal("-1e3")]).fetchone()
+            assert isinstance(ret[0], int)
+    finally:
+        snowflake.connector.paramstyle = original_style
+
+
+def test_decfloat_from_compiler(conn_cnx):
+    # set required decimal precision
+    decimal.getcontext().prec = 38
+    # test both result formats
+    for fmt in ["json", "arrow"]:
+        with conn_cnx(
+            session_parameters={
+                "PYTHON_CONNECTOR_QUERY_RESULT_FORMAT": fmt,
+                "use_cached_result": "false",
+            }
+        ) as cnx:
+            # test endianess
+            ret = cnx.cursor().execute("SELECT 555::decfloat").fetchone()
+            assert isinstance(ret[0], Decimal)
+            assert ret[0] == Decimal("555")
+            # test with decimal separator
+            ret = cnx.cursor().execute("SELECT 123456789.12345678::decfloat").fetchone()
+            assert isinstance(ret[0], Decimal)
+            assert ret[0] == Decimal("123456789.12345678")
+            # test 38 digits
+            ret = (
+                cnx.cursor()
+                .execute("SELECT '12345678901234567890123456789012345678'::decfloat")
+                .fetchone()
+            )
+            assert isinstance(ret[0], Decimal)
+            assert ret[0] == Decimal("12345678901234567890123456789012345678")
+    # test numpy
+    with conn_cnx(numpy=True) as cnx:
+        ret = (
+            cnx.cursor()
+            .execute(
+                "SELECT 1.234::decfloat",
+                None,
+            )
+            .fetchone()
+        )
+        assert isinstance(ret[0], numpy.float64)
+        assert ret[0] == numpy.float64("1.234")