snowflake-connector-python 4.4.0__cp314-cp314-win_amd64.whl

snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowIterator.cpp

@@ -0,0 +1,125 @@
+#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/nanoarrow_cpp/ArrowIterator/CArrowIterator.hpp

@@ -0,0 +1,115 @@
+#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"
+
+#define SF_CHECK_ARROW_RC(arrow_status, 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;                                                         \
+  }
+
+#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