snowflake-connector-python 5.2.0__cp315-cp315-win_amd64.whl

snowflake/connector/_internal/nanoarrow_cpp/ArrowIterator/ArrowTableConverter.hpp

@@ -0,0 +1,63 @@
+#ifndef PC_ARROWTABLECONVERTER_HPP
+#define PC_ARROWTABLECONVERTER_HPP
+
+#include <string>
+
+#include "CArrowIterator.hpp"
+#include "logging.hpp"
+#include "nanoarrow.h"
+#include "nanoarrow.hpp"
+
+namespace sf {
+
+/**
+ * Stateless-ish helper that applies Snowflake-specific type conversions to
+ * Arrow columns in-place.  Both CArrowTableIterator (IPC-based) and
+ * CArrowStreamTableIterator (stream-based) delegate all per-column
+ * conversion to this class so the logic lives in a single place.
+ *
+ * Construction captures the three session-level knobs that affect
+ * conversion behaviour; individual columns are converted via
+ * convertIfNeeded().
+ */
+class ArrowTableConverter {
+ public:
+    ArrowTableConverter(bool number_to_decimal, bool force_microsecond_precision, const std::string& timezone);
+
+  /**
+   * Inspect the Snowflake logical-type metadata on @p columnSchema and
+   * convert @p columnArray in-place when necessary.
+   *
+   * Sets a Python exception on failure (caller should check
+   * py::checkPyError() after the call).
+   */
+  void convertIfNeeded(ArrowSchema* columnSchema, ArrowArrayView* columnArray);
+
+ private:
+  void convertScaledFixedNumberColumn_nanoarrow(ArrowSchemaView* field, ArrowArrayView* columnArray, unsigned int scale);
+
+  void convertScaledFixedNumberColumnToDecimalColumn_nanoarrow(ArrowSchemaView* field, ArrowArrayView* columnArray, unsigned int scale);
+
+  void convertScaledFixedNumberColumnToDoubleColumn_nanoarrow(ArrowSchemaView* field, ArrowArrayView* columnArray, unsigned int scale);
+
+  void convertTimeColumn_nanoarrow(ArrowSchemaView* field, ArrowArrayView* columnArray, int scale);
+
+  void convertIntervalDayTimeColumn_nanoarrow(ArrowSchemaView* field, ArrowArrayView* columnArray, int scale);
+
+  void convertTimestampColumn_nanoarrow(ArrowSchemaView* field, ArrowArrayView* columnArray, int scale, const std::string& timezone = "");
+
+  void convertTimestampTZColumn_nanoarrow(ArrowSchemaView* field, ArrowArrayView* columnArray, int scale, int byteLength, const std::string& timezone);
+
+  template <typename T>
+  double convertScaledFixedNumberToDouble(unsigned int scale, T originalValue);
+
+  const bool m_convert_number_to_decimal;
+  const bool m_force_microsecond_precision;
+  const std::string m_timezone;
+
+  static Logger* logger;
+};
+
+}  // namespace sf
+
+#endif  // PC_ARROWTABLECONVERTER_HPP

snowflake/connector/_internal/nanoarrow_cpp/ArrowIterator/BinaryConverter.cpp

@@ -0,0 +1,18 @@
+#include "BinaryConverter.hpp"
+
+#include <memory>
+
+namespace sf {
+Logger* BinaryConverter::logger = new Logger("snowflake.connector.BinaryConverter");
+
+BinaryConverter::BinaryConverter(ArrowArrayView* array) : m_array(array) {}
+
+PyObject* BinaryConverter::toPyObject(int64_t rowIndex) const {
+  if (ArrowArrayViewIsNull(m_array, rowIndex)) {
+    Py_RETURN_NONE;
+  }
+  ArrowStringView stringView = ArrowArrayViewGetStringUnsafe(m_array, rowIndex);
+  return PyByteArray_FromStringAndSize(stringView.data, stringView.size_bytes);
+}
+
+}  // namespace sf

snowflake/connector/_internal/nanoarrow_cpp/ArrowIterator/BinaryConverter.hpp

@@ -0,0 +1,26 @@
+#ifndef PC_BINARYCONVERTER_HPP
+#define PC_BINARYCONVERTER_HPP
+
+#include <memory>
+
+#include "IColumnConverter.hpp"
+#include "logging.hpp"
+#include "nanoarrow.h"
+
+namespace sf {
+
+class BinaryConverter : public IColumnConverter {
+ public:
+  explicit BinaryConverter(ArrowArrayView* array);
+
+  PyObject* toPyObject(int64_t rowIndex) const override;
+
+ private:
+  ArrowArrayView* m_array;
+
+  static Logger* logger;
+};
+
+}  // namespace sf
+
+#endif  // PC_BINARYCONVERTER_HPP

snowflake/connector/_internal/nanoarrow_cpp/ArrowIterator/BooleanConverter.cpp

@@ -0,0 +1,21 @@
+#include "BooleanConverter.hpp"
+
+#include <memory>
+
+namespace sf {
+
+BooleanConverter::BooleanConverter(ArrowArrayView* array) : m_array(array) {}
+
+PyObject* BooleanConverter::toPyObject(int64_t rowIndex) const {
+  if (ArrowArrayViewIsNull(m_array, rowIndex)) {
+    Py_RETURN_NONE;
+  }
+
+  if (ArrowArrayViewGetIntUnsafe(m_array, rowIndex)) {
+    Py_RETURN_TRUE;
+  } else {
+    Py_RETURN_FALSE;
+  }
+}
+
+}  // namespace sf

snowflake/connector/_internal/nanoarrow_cpp/ArrowIterator/BooleanConverter.hpp

@@ -0,0 +1,23 @@
+#ifndef PC_BOOLEANCONVERTER_HPP
+#define PC_BOOLEANCONVERTER_HPP
+
+#include <memory>
+
+#include "IColumnConverter.hpp"
+#include "nanoarrow.h"
+
+namespace sf {
+
+class BooleanConverter : public IColumnConverter {
+ public:
+  explicit BooleanConverter(ArrowArrayView* array);
+
+  PyObject* toPyObject(int64_t rowIndex) const override;
+
+ private:
+  ArrowArrayView* m_array;
+};
+
+}  // namespace sf
+
+#endif  // PC_BOOLEANCONVERTER_HPP

snowflake/connector/_internal/nanoarrow_cpp/ArrowIterator/CArrowIterator.cpp

@@ -0,0 +1,118 @@
+#include "CArrowIterator.hpp"
+
+#include <memory>
+
+#include "nanoarrow.h"
+#include "nanoarrow_ipc.h"
+
+namespace sf {
+
+const char* const NANOARROW_TYPE_ENUM_STRING[] = {"NANOARROW_TYPE_UNINITIALIZED",
+                                                  "NANOARROW_TYPE_NA",
+                                                  "NANOARROW_TYPE_BOOL",
+                                                  "NANOARROW_TYPE_UINT8",
+                                                  "NANOARROW_TYPE_INT8",
+                                                  "NANOARROW_TYPE_UINT16",
+                                                  "NANOARROW_TYPE_INT16",
+                                                  "NANOARROW_TYPE_UINT32",
+                                                  "NANOARROW_TYPE_INT32",
+                                                  "NANOARROW_TYPE_UINT64",
+                                                  "NANOARROW_TYPE_INT64",
+                                                  "NANOARROW_TYPE_HALF_FLOAT",
+                                                  "NANOARROW_TYPE_FLOAT",
+                                                  "NANOARROW_TYPE_DOUBLE",
+                                                  "NANOARROW_TYPE_STRING",
+                                                  "NANOARROW_TYPE_BINARY",
+                                                  "NANOARROW_TYPE_FIXED_SIZE_BINARY",
+                                                  "NANOARROW_TYPE_DATE32",
+                                                  "NANOARROW_TYPE_DATE64",
+                                                  "NANOARROW_TYPE_TIMESTAMP",
+                                                  "NANOARROW_TYPE_TIME32",
+                                                  "NANOARROW_TYPE_TIME64",
+                                                  "NANOARROW_TYPE_INTERVAL_MONTHS",
+                                                  "NANOARROW_TYPE_INTERVAL_DAY_TIME",
+                                                  "NANOARROW_TYPE_DECIMAL128",
+                                                  "NANOARROW_TYPE_DECIMAL256",
+                                                  "NANOARROW_TYPE_LIST",
+                                                  "NANOARROW_TYPE_STRUCT",
+                                                  "NANOARROW_TYPE_SPARSE_UNION",
+                                                  "NANOARROW_TYPE_DENSE_UNION",
+                                                  "NANOARROW_TYPE_DICTIONARY",
+                                                  "NANOARROW_TYPE_MAP",
+                                                  "NANOARROW_TYPE_EXTENSION",
+                                                  "NANOARROW_TYPE_FIXED_SIZE_LIST",
+                                                  "NANOARROW_TYPE_DURATION",
+                                                  "NANOARROW_TYPE_LARGE_STRING",
+                                                  "NANOARROW_TYPE_LARGE_BINARY",
+                                                  "NANOARROW_TYPE_LARGE_LIST",
+                                                  "NANOARROW_TYPE_INTERVAL_MONTH_DAY_NANO"};
+
+Logger* CArrowIterator::logger = new Logger("snowflake.connector.CArrowIterator");
+
+CArrowIterator::CArrowIterator(char* arrow_bytes, int64_t arrow_bytes_size) {
+  int returnCode = 0;
+  ArrowBuffer input_buffer;
+  ArrowBufferInit(&input_buffer);
+  returnCode = ArrowBufferAppend(&input_buffer, arrow_bytes, arrow_bytes_size);
+  SF_CHECK_ARROW_RC(returnCode, "[Snowflake Exception] error loading arrow bytes, error code: %d",
+                    returnCode);
+  ArrowIpcInputStream input;
+  returnCode = ArrowIpcInputStreamInitBuffer(&input, &input_buffer);
+  SF_CHECK_ARROW_RC(returnCode,
+                    "[Snowflake Exception] error initializing "
+                    "ArrowIpcInputStream, error code: %d",
+                    returnCode);
+  ArrowArrayStream stream;
+  returnCode = ArrowIpcArrayStreamReaderInit(&stream, &input, nullptr);
+  SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(
+      returnCode, stream,
+      "[Snowflake Exception] error initializing ArrowIpcArrayStreamReader, "
+      "error code: %d",
+      returnCode);
+  returnCode = stream.get_schema(&stream, m_ipcArrowSchema.get());
+  SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(
+      returnCode, stream, "[Snowflake Exception] error getting schema from stream, error code: %d",
+      returnCode);
+
+  while (true) {
+    nanoarrow::UniqueArray newUniqueArray;
+    nanoarrow::UniqueArrayView newUniqueArrayView;
+    auto retcode = stream.get_next(&stream, newUniqueArray.get());
+    if (retcode == NANOARROW_OK && newUniqueArray->release != nullptr) {
+      m_ipcArrowArrayVec.push_back(std::move(newUniqueArray));
+
+      ArrowError error;
+      returnCode =
+          ArrowArrayViewInitFromSchema(newUniqueArrayView.get(), m_ipcArrowSchema.get(), &error);
+      SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(
+          returnCode, stream,
+          "[Snowflake Exception] error initializing ArrowArrayView from schema "
+          ": %s, error code: %d",
+          ArrowErrorMessage(&error), returnCode);
+
+      returnCode = ArrowArrayViewSetArray(newUniqueArrayView.get(), newUniqueArray.get(), &error);
+      SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(
+          returnCode, stream,
+          "[Snowflake Exception] error setting ArrowArrayView from array : %s, "
+          "error code: %d",
+          ArrowErrorMessage(&error), returnCode);
+      m_ipcArrowArrayViewVec.push_back(std::move(newUniqueArrayView));
+    } else {
+      SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(
+          retcode, stream,
+          "[Snowflake Exception] error getting schema from stream, error code: "
+          "%d",
+          returnCode);
+      break;
+    }
+  }
+  stream.release(&stream);
+  logger->debug(__FILE__, __func__, __LINE__, "Arrow BatchSize: %d", m_ipcArrowArrayVec.size());
+}
+
+ReturnVal CArrowIterator::checkInitializationStatus() {
+  SF_CHECK_PYTHON_ERR()
+  return ReturnVal(nullptr, nullptr);
+}
+
+}  // namespace sf

snowflake/connector/_internal/nanoarrow_cpp/ArrowIterator/CArrowIterator.hpp

@@ -0,0 +1,116 @@
+#ifndef PC_ARROWITERATOR_HPP
+#define PC_ARROWITERATOR_HPP
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "Python/Common.hpp"
+#include "logging.hpp"
+#include "nanoarrow.hpp"
+
+// logger->error() and PyErr_SetString() both call into the Python/C API. This
+// may run inside a Py_BEGIN_ALLOW_THREADS block (e.g. ArrowTableConverter's
+// per-column conversion loop) as well as ordinary GIL-holding code paths, so
+// guard with py::PyUniqueLock, which re-acquires the GIL only if it isn't
+// already held by this thread.
+#define SF_CHECK_ARROW_RC(arrow_status, format_string, ...)                     \
+  if (arrow_status != NANOARROW_OK) {                                           \
+    std::string errorInfo = Logger::formatString(format_string, ##__VA_ARGS__); \
+    {                                                                           \
+      py::PyUniqueLock gilGuard;                                                \
+      logger->error(__FILE__, __func__, __LINE__, errorInfo.c_str());           \
+      PyErr_SetString(PyExc_Exception, errorInfo.c_str());                      \
+    }                                                                          \
+    return;                                                                     \
+  }
+
+#define SF_CHECK_ARROW_RC_AND_RETURN(arrow_status, ret_val, format_string, ...) \
+  if (arrow_status != NANOARROW_OK) {                                           \
+    std::string errorInfo = Logger::formatString(format_string, ##__VA_ARGS__); \
+    logger->error(__FILE__, __func__, __LINE__, errorInfo.c_str());             \
+    PyErr_SetString(PyExc_Exception, errorInfo.c_str());                        \
+    return ret_val;                                                             \
+  }
+
+#define SF_CHECK_ARROW_RC_AND_RELEASE_ARROW_STREAM(arrow_status, stream, format_string, ...) \
+  if (arrow_status != NANOARROW_OK) {                                                        \
+    std::string errorInfo = std::string(format_string) + std::string(", error info: ") +     \
+                            std::string(stream.get_last_error(&stream));                     \
+    std::string fullErrorInfo = Logger::formatString(errorInfo.c_str(), ##__VA_ARGS__);      \
+    logger->error(__FILE__, __func__, __LINE__, fullErrorInfo.c_str());                      \
+    PyErr_SetString(PyExc_Exception, fullErrorInfo.c_str());                                 \
+    stream.release(&stream);                                                                 \
+    return;                                                                                  \
+  }
+
+#define SF_CHECK_PYTHON_ERR()                              \
+  if (py::checkPyError()) {                                \
+    PyObject *type, *val, *traceback;                      \
+    PyErr_Fetch(&type, &val, &traceback);                  \
+    PyErr_Clear();                                         \
+    m_currentPyException.reset(val);                       \
+                                                           \
+    Py_XDECREF(type);                                      \
+    Py_XDECREF(traceback);                                 \
+                                                           \
+    return ReturnVal(nullptr, m_currentPyException.get()); \
+  }
+
+namespace sf {
+
+extern const char* const NANOARROW_TYPE_ENUM_STRING[];
+
+/**
+ * A simple struct to contain return data back cython.
+ * PyObject would be nullptr if failed and cause string will be populated
+ *
+ * Note that `ReturnVal` does not own these pointers, so they should
+ * not be decref'ed by the receiver.
+ */
+class ReturnVal {
+ public:
+  ReturnVal() : successObj(nullptr), exception(nullptr) {}
+
+  ReturnVal(PyObject* obj, PyObject* except) : successObj(obj), exception(except) {}
+
+  PyObject* successObj;
+
+  PyObject* exception;
+};
+
+/**
+ * Arrow base iterator implementation in C++.
+ */
+
+class CArrowIterator {
+ public:
+  CArrowIterator(char* arrow_bytes, int64_t arrow_bytes_size);
+
+  virtual ~CArrowIterator() = default;
+
+  /**
+   * @return a python object which might be current row or an Arrow Table
+   */
+  virtual ReturnVal next() = 0;
+  virtual std::vector<uintptr_t> getArrowArrayPtrs() { return {}; };
+  virtual std::vector<uintptr_t> getArrowSchemaPtrs() { return {}; };
+
+  /** check whether initialization succeeded or encountered error */
+  ReturnVal checkInitializationStatus();
+
+ protected:
+  static Logger* logger;
+
+  /** nanoarrow data */
+  std::vector<nanoarrow::UniqueArray> m_ipcArrowArrayVec;
+  std::vector<nanoarrow::UniqueArrayView> m_ipcArrowArrayViewVec;
+  nanoarrow::UniqueSchema m_ipcArrowSchema;
+
+  /** pointer to the current python exception object */
+  py::UniqueRef m_currentPyException;
+};
+}  // namespace sf
+
+#endif  // PC_ARROWITERATOR_HPP