SNOW-773928: nanoarrow decimal support (#1538)

Author: sfc-gh-aling

src/snowflake/connector/cpp/ArrowIterator/CArrowTableIterator.cpp

@@ -312,7 +312,7 @@ void CArrowTableIterator::convertScaledFixedNumberColumnToDecimalColumn_nanoarro
   // create new schema
   ArrowSchemaInit(newSchema);
   newSchema->flags &= (field->schema->flags & ARROW_FLAG_NULLABLE); // map to nullable()
-  ArrowSchemaSetType(newSchema, NANOARROW_TYPE_DECIMAL128);  // map to arrow:float64()
+  ArrowSchemaSetTypeDecimal(newSchema, NANOARROW_TYPE_DECIMAL128, 38, scale);
   ArrowSchemaSetName(newSchema, field->schema->name);
 
   ArrowError error;
@@ -325,14 +325,17 @@ void CArrowTableIterator::convertScaledFixedNumberColumnToDecimalColumn_nanoarro
     logger->error(__FILE__, __func__, __LINE__, errorInfo.c_str());
     PyErr_SetString(PyExc_Exception, errorInfo.c_str());
   }
-
+  ArrowArrayStartAppending(newArray);
   for(int64_t rowIdx = 0; rowIdx < columnArray->array->length; rowIdx++)
   {
     if(ArrowArrayViewIsNull(columnArray, rowIdx)) {
         ArrowArrayAppendNull(newArray, 1);
     } else {
         auto originalVal = ArrowArrayViewGetIntUnsafe(columnArray, rowIdx);
-        // TODO: nanoarrow is missing appending a decimal value to array
+        std::shared_ptr<ArrowDecimal> arrowDecimal = std::make_shared<ArrowDecimal>();
+        ArrowDecimalInit(arrowDecimal.get(), 128, 38, scale);
+        ArrowDecimalSetInt(arrowDecimal.get(), originalVal);
+        ArrowArrayAppendDecimal(newArray, arrowDecimal.get());
     }
   }
   ArrowArrayFinishBuildingDefault(newArray, &error);

src/snowflake/connector/cpp/ArrowIterator/nanoarrow.h

@@ -204,6 +204,13 @@ typedef int ArrowErrorCode;
 #define NANOARROW_RETURN_NOT_OK(EXPR) \
   _NANOARROW_RETURN_NOT_OK_IMPL(_NANOARROW_MAKE_NAME(errno_status_, __COUNTER__), EXPR)
 
+static char _ArrowIsLittleEndian(void) {
+  uint32_t check = 1;
+  char first_byte;
+  memcpy(&first_byte, &check, sizeof(char));
+  return first_byte;
+}
+
 /// \brief Arrow type enumerator
 /// \ingroup nanoarrow-utils
 ///
@@ -586,6 +593,90 @@ struct ArrowArrayPrivateData {
   int8_t union_type_id_is_child_index;
 };
 
+/// \brief A representation of a fixed-precision decimal number
+/// \ingroup nanoarrow-utils
+///
+/// This structure should be initialized with ArrowDecimalInit() once and
+/// values set using ArrowDecimalSetInt(), ArrowDecimalSetBytes128(),
+/// or ArrowDecimalSetBytes256().
+struct ArrowDecimal {
+  /// \brief An array of 64-bit integers of n_words length defined in native-endian order
+  uint64_t words[4];
+
+  /// \brief The number of significant digits this decimal number can represent
+  int32_t precision;
+
+  /// \brief The number of digits after the decimal point. This can be negative.
+  int32_t scale;
+
+  /// \brief The number of words in the words array
+  int n_words;
+
+  /// \brief Cached value used by the implementation
+  int high_word_index;
+
+  /// \brief Cached value used by the implementation
+  int low_word_index;
+};
+
+/// \brief Initialize a decimal with a given set of type parameters
+/// \ingroup nanoarrow-utils
+static inline void ArrowDecimalInit(struct ArrowDecimal* decimal, int32_t bitwidth,
+                                    int32_t precision, int32_t scale) {
+  memset(decimal->words, 0, sizeof(decimal->words));
+  decimal->precision = precision;
+  decimal->scale = scale;
+  decimal->n_words = bitwidth / 8 / sizeof(uint64_t);
+
+  if (_ArrowIsLittleEndian()) {
+    decimal->low_word_index = 0;
+    decimal->high_word_index = decimal->n_words - 1;
+  } else {
+    decimal->low_word_index = decimal->n_words - 1;
+    decimal->high_word_index = 0;
+  }
+}
+
+/// \brief Get a signed integer value of a sufficiently small ArrowDecimal
+///
+/// This does not check if the decimal's precision sufficiently small to fit
+/// within the signed 64-bit integer range (A precision less than or equal
+/// to 18 is sufficiently small).
+static inline int64_t ArrowDecimalGetIntUnsafe(struct ArrowDecimal* decimal) {
+  return (int64_t)decimal->words[decimal->low_word_index];
+}
+
+/// \brief Copy the bytes of this decimal into a sufficiently large buffer
+/// \ingroup nanoarrow-utils
+static inline void ArrowDecimalGetBytes(struct ArrowDecimal* decimal, uint8_t* out) {
+  memcpy(out, decimal->words, decimal->n_words * sizeof(uint64_t));
+}
+
+/// \brief Returns 1 if the value represented by decimal is >= 0 or -1 otherwise
+/// \ingroup nanoarrow-utils
+static inline int64_t ArrowDecimalSign(struct ArrowDecimal* decimal) {
+  return 1 | ((int64_t)(decimal->words[decimal->high_word_index]) >> 63);
+}
+
+/// \brief Sets the integer value of this decimal
+/// \ingroup nanoarrow-utils
+static inline void ArrowDecimalSetInt(struct ArrowDecimal* decimal, int64_t value) {
+  if (value < 0) {
+    memset(decimal->words, 0xff, decimal->n_words * sizeof(uint64_t));
+  } else {
+    memset(decimal->words, 0, decimal->n_words * sizeof(uint64_t));
+  }
+
+  decimal->words[decimal->low_word_index] = value;
+}
+
+/// \brief Copy bytes from a buffer into this decimal
+/// \ingroup nanoarrow-utils
+static inline void ArrowDecimalSetBytes(struct ArrowDecimal* decimal,
+                                        const uint8_t* value) {
+  memcpy(decimal->words, value, decimal->n_words * sizeof(uint64_t));
+}
+
 #ifdef __cplusplus
 }
 #endif
@@ -1417,12 +1508,20 @@ static inline ArrowErrorCode ArrowArrayAppendBytes(struct ArrowArray* array,
                                                    struct ArrowBufferView value);
 
 /// \brief Append a string value to an array
+///
 /// Returns NANOARROW_OK if value can be exactly represented by
 /// the underlying storage type or EINVAL otherwise (e.g.,
 /// the underlying array is not a string or large string array).
 static inline ArrowErrorCode ArrowArrayAppendString(struct ArrowArray* array,
                                                     struct ArrowStringView value);
 
+/// \brief Append a decimal value to an array
+///
+/// Returns NANOARROW_OK if array is a decimal array with the appropriate
+/// bitwidth or EINVAL otherwise.
+static inline ArrowErrorCode ArrowArrayAppendDecimal(struct ArrowArray* array,
+                                                     struct ArrowDecimal* value);
+
 /// \brief Finish a nested array element
 ///
 /// Appends a non-null element to the array based on the first child's current
@@ -1559,6 +1658,14 @@ static inline struct ArrowStringView ArrowArrayViewGetStringUnsafe(
 static inline struct ArrowBufferView ArrowArrayViewGetBytesUnsafe(
     struct ArrowArrayView* array_view, int64_t i);
 
+/// \brief Get an element in an ArrowArrayView as an ArrowDecimal
+///
+/// This function does not check for null values. The out parameter must
+/// be initialized with ArrowDecimalInit() with the proper parameters for this
+/// type before calling this for the first time.
+static inline void ArrowArrayViewGetDecimalUnsafe(struct ArrowArrayView* array_view,
+                                                  int64_t i, struct ArrowDecimal* out);
+
 /// @}
 
 /// \defgroup nanoarrow-basic-array-stream Basic ArrowArrayStream implementation
@@ -2622,6 +2729,41 @@ static inline ArrowErrorCode ArrowArrayAppendString(struct ArrowArray* array,
   }
 }
 
+static inline ArrowErrorCode ArrowArrayAppendDecimal(struct ArrowArray* array,
+                                                     struct ArrowDecimal* value) {
+  struct ArrowArrayPrivateData* private_data =
+      (struct ArrowArrayPrivateData*)array->private_data;
+  struct ArrowBuffer* data_buffer = ArrowArrayBuffer(array, 1);
+
+  switch (private_data->storage_type) {
+    case NANOARROW_TYPE_DECIMAL128:
+      if (value->n_words != 2) {
+        return EINVAL;
+      } else {
+        NANOARROW_RETURN_NOT_OK(
+            ArrowBufferAppend(data_buffer, value->words, 2 * sizeof(uint64_t)));
+        break;
+      }
+    case NANOARROW_TYPE_DECIMAL256:
+      if (value->n_words != 4) {
+        return EINVAL;
+      } else {
+        NANOARROW_RETURN_NOT_OK(
+            ArrowBufferAppend(data_buffer, value->words, 4 * sizeof(uint64_t)));
+        break;
+      }
+    default:
+      return EINVAL;
+  }
+
+  if (private_data->bitmap.buffer.data != NULL) {
+    NANOARROW_RETURN_NOT_OK(ArrowBitmapAppend(ArrowArrayValidityBitmap(array), 1, 1));
+  }
+
+  array->length++;
+  return NANOARROW_OK;
+}
+
 static inline ArrowErrorCode ArrowArrayFinishElement(struct ArrowArray* array) {
   struct ArrowArrayPrivateData* private_data =
       (struct ArrowArrayPrivateData*)array->private_data;
@@ -2931,6 +3073,23 @@ static inline struct ArrowBufferView ArrowArrayViewGetBytesUnsafe(
   return view;
 }
 
+static inline void ArrowArrayViewGetDecimalUnsafe(struct ArrowArrayView* array_view,
+                                                  int64_t i, struct ArrowDecimal* out) {
+  i += array_view->array->offset;
+  const uint8_t* data_view = array_view->buffer_views[1].data.as_uint8;
+  switch (array_view->storage_type) {
+    case NANOARROW_TYPE_DECIMAL128:
+      ArrowDecimalSetBytes(out, data_view + (i * 16));
+      break;
+    case NANOARROW_TYPE_DECIMAL256:
+      ArrowDecimalSetBytes(out, data_view + (i * 32));
+      break;
+    default:
+      memset(out->words, 0, sizeof(out->words));
+      break;
+  }
+}
+
 #ifdef __cplusplus
 }
 #endif

src/snowflake/connector/cpp/ArrowIterator/nanoarrow.hpp

@@ -1,3 +1,7 @@
+//
+// Copyright (c) 2012-2021 Snowflake Computing Inc. All rights reserved.
+//
+
 // Licensed to the Apache Software Foundation (ASF) under one
 // or more contributor license agreements.  See the NOTICE file
 // distributed with this work for additional information
@@ -278,10 +282,10 @@ class VectorArrayStream : public EmptyArrayStream {
 
  protected:
   VectorArrayStream(struct ArrowSchema* schema, std::vector<UniqueArray> arrays)
-      : EmptyArrayStream(schema), offset_(0), arrays_(std::move(arrays)) {}
+      : EmptyArrayStream(schema), arrays_(std::move(arrays)), offset_(0) {}
 
   int get_next(struct ArrowArray* array) {
-    if (offset_ < arrays_.size()) {
+    if (offset_ < static_cast<int64_t>(arrays_.size())) {
       arrays_[offset_++].move(array);
     } else {
       array->release = nullptr;

test/integ/pandas/test_arrow_pandas.py

@@ -936,7 +936,6 @@ def test_query_resultscan_combos(conn_cnx, query_format, resultscan_format):
     ],
 )
 def test_number_fetchall_retrieve_type(conn_cnx, use_decimal, expected):
-    pytest.skip("missing decimal")
     with conn_cnx(arrow_number_to_decimal=use_decimal) as con:
         with con.cursor() as cur:
             cur.execute("SELECT 12345600.87654301::NUMBER(18, 8) a")
@@ -956,7 +955,6 @@ def test_number_fetchall_retrieve_type(conn_cnx, use_decimal, expected):
     ],
 )
 def test_number_fetchbatches_retrieve_type(conn_cnx, use_decimal: bool, expected: type):
-    pytest.skip("missing decimal")
     with conn_cnx(arrow_number_to_decimal=use_decimal) as con:
         with con.cursor() as cur:
             cur.execute("SELECT 12345600.87654301::NUMBER(18, 8) a")
@@ -1249,7 +1247,6 @@ def test_timestamp_tz(conn_cnx):
 
 
 def test_arrow_number_to_decimal(conn_cnx):
-    pytest.skip("missing decimal")
     with conn_cnx(
         session_parameters={
             PARAMETER_PYTHON_CONNECTOR_QUERY_RESULT_FORMAT: "arrow_force"