SNOW-950840 Add support for the vector data type (#1804)

Author: sfc-gh-rpanchapakesan

.gitignore

@@ -124,3 +124,4 @@ core.*
 
 # Compiled Cython
 src/snowflake/connector/arrow_iterator.cpp
+src/snowflake/connector/nanoarrow_cpp/ArrowIterator/nanoarrow_arrow_iterator.cpp

DESCRIPTION.md

@@ -9,6 +9,10 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 # Release Notes
 
 
+- v3.6.0(TBD)
+
+  - Added support for Vector types
+
 - v3.5.0(November 13,2023)
 
   - Version 3.5.0 is the snowflake-connector-python purely built upon apache arrow-nanoarrow project.

setup.py

@@ -107,6 +107,9 @@ def build_extension(self, ext):
                         NANOARROW_ARROW_ITERATOR_SRC_DIR, "DecimalConverter.cpp"
                     ),
                     os.path.join(NANOARROW_ARROW_ITERATOR_SRC_DIR, "DateConverter.cpp"),
+                    os.path.join(
+                        NANOARROW_ARROW_ITERATOR_SRC_DIR, "FixedSizeListConverter.cpp"
+                    ),
                     os.path.join(
                         NANOARROW_ARROW_ITERATOR_SRC_DIR, "FloatConverter.cpp"
                     ),

src/snowflake/connector/constants.py

@@ -95,6 +95,14 @@ class FieldType(NamedTuple):
     FieldType(
         name="GEOMETRY", dbapi_type=[DBAPI_TYPE_STRING], pa_type=lambda: pa.string()
     ),
+    FieldType(
+        # TODO(SNOW-969160): While pa.binary() results in the correct pandas column
+        # type being generated, it should be switched to pa.list_(...) once parsing
+        # for the new result metadata fields is added.
+        name="VECTOR",
+        dbapi_type=[DBAPI_TYPE_BINARY],
+        pa_type=lambda: pa.binary(),
+    ),
 )
 
 FIELD_NAME_TO_ID: DefaultDict[Any, int] = defaultdict(int)

src/snowflake/connector/converter.py

@@ -7,6 +7,7 @@
 
 import binascii
 import decimal
+import json
 import time
 from datetime import date, datetime
 from datetime import time as dt_t
@@ -320,6 +321,9 @@ def _VARIANT_to_python(self, _: dict[str, Any]) -> Any | None:
 
     _ARRAY_to_python = _VARIANT_to_python
 
+    def _VECTOR_to_python(self, ctx: dict[str, Any]) -> Callable:
+        return lambda v: json.loads(v)
+
     def _BOOLEAN_to_python(
         self, ctx: dict[str, str | None] | dict[str, str]
     ) -> Callable:

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

@@ -11,6 +11,7 @@
 #include "BinaryConverter.hpp"
 #include "BooleanConverter.hpp"
 #include "DateConverter.hpp"
+#include "FixedSizeListConverter.hpp"
 #include "TimeStampConverter.hpp"
 #include "TimeConverter.hpp"
 #include <memory>
@@ -438,6 +439,12 @@ void CArrowChunkIterator::initColumnConverters()
         break;
       }
 
+      case SnowflakeType::Type::VECTOR:
+      {
+        m_currentBatchConverters.push_back(std::make_shared<sf::FixedSizeListConverter>(array));
+        break;
+      }
+
       default:
       {
         std::string errorInfo = Logger::formatString(

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

@@ -0,0 +1,88 @@
+//
+// Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+//
+
+#include "FixedSizeListConverter.hpp"
+
+namespace sf
+{
+Logger* FixedSizeListConverter::logger =
+    new Logger("snowflake.connector.FixedSizeListConverter");
+
+FixedSizeListConverter::FixedSizeListConverter(ArrowArrayView* array)
+: m_array(array)
+{
+}
+
+void FixedSizeListConverter::generateError(const std::string& msg) const
+{
+  logger->error(__FILE__, __func__, __LINE__, msg.c_str());
+  PyErr_SetString(PyExc_Exception, msg.c_str());
+}
+
+PyObject* FixedSizeListConverter::toPyObject(int64_t rowIndex) const
+{
+  if (ArrowArrayViewIsNull(m_array, rowIndex))
+  {
+    Py_RETURN_NONE;
+  }
+
+  if (m_array->n_children != 1)
+  {
+    std::string errorInfo = Logger::formatString(
+        "[Snowflake Exception] invalid arrow element schema for fixed size "
+        "list: got (%d) "
+        "children",
+        m_array->n_children);
+    this->generateError(errorInfo);
+    return nullptr;
+  }
+
+  // m_array->length represents the number of fixed size lists in the array
+  // m_array->children[0] has a buffer view that contains the actual data of
+  // each list, back-to-back m_array->children[0]->length represents the sum of
+  // the lengths of the fixed size lists in the array.
+
+  ArrowArrayView* elements = m_array->children[0];
+  const auto fixedSizeArrayLength = elements->length / m_array->length;
+  PyObject* list = PyList_New(fixedSizeArrayLength);
+
+  const int64_t startIndexWithoutOffset = rowIndex * fixedSizeArrayLength;
+  for (int64_t i = 0; i < fixedSizeArrayLength; ++i)
+  {
+    const auto bufferIndexWithoutOffset = startIndexWithoutOffset + i;
+    // Currently, the backend only sends back INT32 and FLOAT32, but the
+    // remaining types are enumerated for future use.
+    switch (elements->storage_type)
+    {
+      case NANOARROW_TYPE_INT8:
+      case NANOARROW_TYPE_INT16:
+      case NANOARROW_TYPE_INT32:
+      case NANOARROW_TYPE_INT64:
+      {
+        const auto value =
+            ArrowArrayViewGetIntUnsafe(elements, bufferIndexWithoutOffset);
+        PyList_SetItem(list, i, PyLong_FromLongLong(value));
+      } break;
+      case NANOARROW_TYPE_HALF_FLOAT:
+      case NANOARROW_TYPE_FLOAT:
+      case NANOARROW_TYPE_DOUBLE:
+      {
+        const auto value =
+            ArrowArrayViewGetDoubleUnsafe(elements, bufferIndexWithoutOffset);
+        PyList_SetItem(list, i, PyFloat_FromDouble(value));
+      } break;
+      default:
+        std::string errorInfo = Logger::formatString(
+            "[Snowflake Exception] invalid arrow element type for fixed size "
+            "list: got (%s)",
+            ArrowTypeString(elements->storage_type));
+        this->generateError(errorInfo);
+        return nullptr;
+    }
+  }
+
+  return list;
+}
+
+}  // namespace sf

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

@@ -0,0 +1,34 @@
+//
+// Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
+//
+
+#ifndef PC_FIXEDSIZELISTCONVERTER_HPP
+#define PC_FIXEDSIZELISTCONVERTER_HPP
+
+#include <memory>
+
+#include "IColumnConverter.hpp"
+#include "logging.hpp"
+#include "nanoarrow.h"
+#include "nanoarrow.hpp"
+
+namespace sf
+{
+
+class FixedSizeListConverter : public IColumnConverter
+{
+public:
+  explicit FixedSizeListConverter(ArrowArrayView* array);
+  PyObject* toPyObject(int64_t rowIndex) const override;
+
+private:
+  void generateError(const std::string& msg) const;
+
+  ArrowArrayView* m_array;
+
+  static Logger* logger;
+};
+
+}  // namespace sf
+
+#endif  // PC_FIXEDSIZELISTCONVERTER_HPP

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

@@ -29,6 +29,7 @@ std::unordered_map<std::string, SnowflakeType::Type>
         {"TIMESTAMP_LTZ", SnowflakeType::Type::TIMESTAMP_LTZ},
         {"TIMESTAMP_NTZ", SnowflakeType::Type::TIMESTAMP_NTZ},
         {"TIMESTAMP_TZ", SnowflakeType::Type::TIMESTAMP_TZ},
-        {"VARIANT", SnowflakeType::Type::VARIANT}};
+        {"VARIANT", SnowflakeType::Type::VARIANT},
+        {"VECTOR", SnowflakeType::Type::VECTOR}};
 
 }  // namespace sf

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

@@ -32,7 +32,8 @@ class SnowflakeType
     TIMESTAMP_LTZ = 12,
     TIMESTAMP_NTZ = 13,
     TIMESTAMP_TZ = 14,
-    VARIANT = 15
+    VARIANT = 15,
+    VECTOR = 16
   };
 
   static SnowflakeType::Type snowflakeTypeFromString(std::string str)