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

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

@@ -0,0 +1,275 @@
+#include "CArrowStreamTableIterator.hpp"
+
+#include <exception>
+#include <memory>
+#include <string>
+
+#include "Python/Common.hpp"
+
+namespace sf {
+
+Logger* CArrowStreamTableIterator::logger =
+    new Logger("snowflake.connector.CArrowStreamTableIterator");
+
+namespace {
+  void releaseStream(ArrowArrayStream* s) {
+    if (s && s->release) s->release(s);
+  }
+
+  /**
+   * Invokes @p convert, translating any C++ exception into a Python
+   * exception. The Cython declarations of the methods below carry no
+   * `except +`, so an escaping C++ exception reaches a caller that cannot
+   * handle it and terminates the process. py::setPyError() acquires the GIL
+   * on its own, so this is usable inside a Py_BEGIN_ALLOW_THREADS block.
+   */
+  template <typename Fn>
+  void runTranslatingCxxExceptions(Fn&& convert) {
+    try {
+      convert();
+    } catch (const std::exception& e) {
+      py::setPyError(PyExc_Exception, e.what());
+    } catch (...) {
+      py::setPyError(PyExc_Exception,
+                     "[Snowflake Exception] unknown error while converting Arrow columns");
+    }
+  }
+}  // namespace
+
+// ---------------------------------------------------------------------------
+// Factory / constructor
+// ---------------------------------------------------------------------------
+
+std::unique_ptr<CArrowStreamTableIterator> CArrowStreamTableIterator::from_stream(int64_t stream_ptr, PyObject* context, bool number_to_decimal, bool force_microsecond_precision) {
+  auto* stream = reinterpret_cast<ArrowArrayStream*>(stream_ptr);
+  if (!stream || !stream->release) {
+    PyErr_SetString(PyExc_Exception, "[Snowflake Exception] Invalid ArrowArrayStream pointer");
+    return nullptr;
+  }
+
+  ArrowSchema schema{};
+  int rc = stream->get_schema(stream, &schema);
+  if (rc != 0) {
+    const char* msg = stream->get_last_error(stream);
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error getting schema: %s, error code: %d", msg ? msg : "unknown", rc);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_Exception, err.c_str());
+    return nullptr;
+  }
+
+  return std::unique_ptr<CArrowStreamTableIterator>(
+    new CArrowStreamTableIterator(stream, &schema, context, number_to_decimal, force_microsecond_precision));
+}
+
+CArrowStreamTableIterator::CArrowStreamTableIterator(ArrowArrayStream* stream, ArrowSchema* schema, PyObject* context, bool number_to_decimal, bool force_microsecond_precision)
+    : m_stream(stream, releaseStream),
+      m_columnCount(schema->n_children),
+      m_context(context) {
+  ArrowSchemaMove(schema, m_schema.get());
+
+  char* timezone = nullptr;
+  py::UniqueRef tz(PyObject_GetAttrString(m_context, "_timezone"));
+  PyArg_Parse(tz.get(), "z", &timezone);
+
+  m_converter = std::unique_ptr<ArrowTableConverter>(
+    new ArrowTableConverter(number_to_decimal, force_microsecond_precision, timezone ? std::string(timezone) : std::string()));
+}
+
+// ---------------------------------------------------------------------------
+// Export accessors
+// ---------------------------------------------------------------------------
+
+uintptr_t CArrowStreamTableIterator::getArrowArrayPtr() {
+  return reinterpret_cast<uintptr_t>(m_exportArray.get());
+}
+
+uintptr_t CArrowStreamTableIterator::getArrowSchemaPtr() {
+  return reinterpret_cast<uintptr_t>(m_exportSchema.get());
+}
+
+uintptr_t CArrowStreamTableIterator::getConvertedSchemaPtr() {
+  m_streamSchemaExport.reset();
+  int rc = ArrowSchemaDeepCopy(m_schema.get(), m_streamSchemaExport.get());
+  if (rc != NANOARROW_OK) {
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error deep-copying stream schema, error code: %d", rc);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_RuntimeError, err.c_str());
+    return 0;
+  }
+
+  ArrowError error;
+
+  nanoarrow::UniqueArray emptyArray;
+  rc = ArrowArrayInitFromSchema(emptyArray.get(), m_schema.get(), &error);
+  if (rc != NANOARROW_OK) {
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error initializing empty array for schema conversion: %s, error code: %d",
+        ArrowErrorMessage(&error), rc);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_RuntimeError, err.c_str());
+    return 0;
+  }
+  rc = ArrowArrayFinishBuildingDefault(emptyArray.get(), &error);
+  if (rc != NANOARROW_OK) {
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error finishing empty array: %s, error code: %d",
+        ArrowErrorMessage(&error), rc);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_RuntimeError, err.c_str());
+    return 0;
+  }
+
+  nanoarrow::UniqueArrayView emptyView;
+  rc = ArrowArrayViewInitFromSchema(emptyView.get(), m_schema.get(), &error);
+  if (rc != NANOARROW_OK) {
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error initializing empty array view: %s, error code: %d",
+        ArrowErrorMessage(&error), rc);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_RuntimeError, err.c_str());
+    return 0;
+  }
+  rc = ArrowArrayViewSetArray(emptyView.get(), emptyArray.get(), &error);
+  if (rc != NANOARROW_OK) {
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error setting empty array view: %s, error code: %d",
+        ArrowErrorMessage(&error), rc);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_RuntimeError, err.c_str());
+    return 0;
+  }
+
+  runTranslatingCxxExceptions([this, &emptyView] {
+    for (int64_t col = 0; col < m_columnCount; col++) {
+      m_converter->convertIfNeeded(m_streamSchemaExport->children[col], emptyView->children[col]);
+      if (py::checkPyError()) return;
+    }
+  });
+  if (py::checkPyError()) return 0;
+
+  return reinterpret_cast<uintptr_t>(m_streamSchemaExport.get());
+}
+
+// ---------------------------------------------------------------------------
+// Stream loading
+// ---------------------------------------------------------------------------
+
+bool CArrowStreamTableIterator::loadNextBatch() {
+  if (m_streamExhausted) return false;
+
+  m_currentArray.reset();
+  m_currentArrayView.reset();
+
+  int returnCode;
+  ArrowArrayStream* stream = m_stream.get();
+  {
+    Py_BEGIN_ALLOW_THREADS
+    returnCode = stream->get_next(stream, m_currentArray.get());
+    Py_END_ALLOW_THREADS
+  }
+
+  if (returnCode != NANOARROW_OK) {
+    const char* msg = stream->get_last_error(stream);
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error getting next batch: %s, error code: %d", msg ? msg : "unknown", returnCode);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_Exception, err.c_str());
+    return false;
+  }
+
+  if (m_currentArray->release == nullptr) {
+    m_streamExhausted = true;
+    return false;
+  }
+
+  if (m_currentArray->length == 0) {
+    return loadNextBatch();
+  }
+
+  ArrowError error;
+  returnCode = ArrowArrayViewInitFromSchema(m_currentArrayView.get(), m_schema.get(), &error);
+  if (returnCode != NANOARROW_OK) {
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error initializing ArrowArrayView: %s, error code: %d",
+        ArrowErrorMessage(&error), returnCode);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_Exception, err.c_str());
+    return false;
+  }
+
+  returnCode = ArrowArrayViewSetArray(m_currentArrayView.get(), m_currentArray.get(), &error);
+  if (returnCode != NANOARROW_OK) {
+    std::string err = Logger::formatString(
+        "[Snowflake Exception] error setting ArrowArrayView: %s, error code: %d",
+        ArrowErrorMessage(&error), returnCode);
+    logger->error(__FILE__, __func__, __LINE__, err.c_str());
+    PyErr_SetString(PyExc_Exception, err.c_str());
+    return false;
+  }
+
+  return true;
+}
+
+// ---------------------------------------------------------------------------
+// convertBatch()
+// ---------------------------------------------------------------------------
+
+void CArrowStreamTableIterator::convertBatch() {
+  m_exportArray.reset();
+  m_exportSchema.reset();
+
+  int rc = ArrowSchemaDeepCopy(m_schema.get(), m_exportSchema.get());
+  SF_CHECK_ARROW_RC(rc, "[Snowflake Exception] error deep-copying schema, error code: %d", rc);
+
+  {
+    Py_BEGIN_ALLOW_THREADS
+    runTranslatingCxxExceptions([this] {
+      for (int64_t col = 0; col < m_columnCount; col++) {
+        m_converter->convertIfNeeded(m_exportSchema->children[col],
+                                     m_currentArrayView->children[col]);
+      }
+    });
+    Py_END_ALLOW_THREADS
+  }
+  // GIL is held again - surface any error set via py::setPyError
+  if (py::checkPyError()) return;
+
+  ArrowArrayMove(m_currentArray.get(), m_exportArray.get());
+}
+
+// ---------------------------------------------------------------------------
+// next()
+// ---------------------------------------------------------------------------
+
+ReturnVal CArrowStreamTableIterator::next() {
+  if (!loadNextBatch()) {
+    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());
+    }
+    return ReturnVal(Py_None, nullptr);
+  }
+
+  convertBatch();
+
+  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());
+  }
+
+  return ReturnVal(Py_True, nullptr);
+}
+
+}  // namespace sf

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

@@ -0,0 +1,168 @@
+#ifndef PC_CARROWSTREAMTABLEITERATOR_HPP
+#define PC_CARROWSTREAMTABLEITERATOR_HPP
+
+#include <memory>
+
+#include "ArrowTableConverter.hpp"
+#include "CArrowIterator.hpp"
+#include "Python/Common.hpp"
+#include "logging.hpp"
+#include "nanoarrow.h"
+#include "nanoarrow.hpp"
+
+namespace sf {
+
+/**
+ * Stream-based Arrow iterator that reads batches from an ArrowArrayStream
+ * and exposes each batch as a (ArrowArray, ArrowSchema) pair suitable for
+ * import into PyArrow via RecordBatch._import_from_c().
+ *
+ * Snowflake-specific type conversions (e.g. scaled NUMBER -> float64,
+ * struct-encoded TIMESTAMP -> Arrow timestamp) are applied to every batch
+ * through the shared ArrowTableConverter.
+ *
+ * Lifecycle (driven from the Cython wrapper ArrowDummyIterator):
+ *   1. from_stream()   – validate stream, read schema, construct instance
+ *   2. next()          – load next batch, convert, expose via export slots
+ *   3. getArrowArrayPtr() / getArrowSchemaPtr()
+ *                      – caller passes these to _import_from_c() which
+ *                        consumes them (sets release to nullptr)
+ *   4. repeat 2-3 until next() signals exhaustion (Py_None)
+ */
+class CArrowStreamTableIterator {
+ public:
+  /**
+   * Factory method.  Validates the stream pointer, reads the schema, and
+   * returns a ready-to-iterate instance.
+   *
+   * @param stream_ptr  Pointer to ArrowArrayStream as int64 (FFI from Python)
+   * @param context     Python ArrowConverterContext carrying session timezone
+   * @param number_to_decimal
+   *   If true, scaled FIXED columns become Decimal128; otherwise float64
+   * @param force_microsecond_precision
+   *   If true, all TIMESTAMP columns are forced to microsecond precision,
+   *   ensuring consistent schema across batches whose data may span
+   *   outside the nanosecond-representable range (1677-2262)
+   * @return Unique pointer to the iterator, or nullptr with a Python
+   *         exception set on failure
+   */
+  static std::unique_ptr<CArrowStreamTableIterator> from_stream(
+      int64_t stream_ptr, PyObject* context, bool number_to_decimal,
+      bool force_microsecond_precision);
+
+  /**
+   * Advance to the next batch.
+   *
+   * @return ReturnVal with:
+   *   - successObj = Py_True, exception = nullptr   → batch ready; call
+   *     getArrowArrayPtr() / getArrowSchemaPtr() to retrieve it
+   *   - successObj = Py_None, exception = nullptr   → stream exhausted
+   *   - successObj = nullptr, exception != nullptr   → error occurred
+   */
+  ReturnVal next();
+
+  /**
+   * Raw pointer (as uintptr_t) to the converted ArrowArray for the
+   * current batch.  Valid only immediately after next() returns Py_True.
+   * Ownership transfers to the caller via _import_from_c().
+   */
+  uintptr_t getArrowArrayPtr();
+
+  /**
+   * Raw pointer (as uintptr_t) to the deep-copied ArrowSchema for the
+   * current batch.  Same ownership semantics as getArrowArrayPtr().
+   */
+  uintptr_t getArrowSchemaPtr();
+
+  /**
+   * Raw pointer (as uintptr_t) to the stream schema after applying
+   * Snowflake type conversions (e.g. struct → timestamp, int → float64).
+   * Unlike getArrowSchemaPtr(), this is always valid—even when no batch
+   * has been loaded—making it safe to call for empty result sets
+   * (e.g. to build an empty PyArrow table with the correct column types).
+   *
+   * Internally synthesises a zero-row batch and runs it through
+   * ArrowTableConverter so the schema matches what non-empty results
+   * would produce.
+   *
+   * Ownership transfers to the caller via _import_from_c().
+   */
+  uintptr_t getConvertedSchemaPtr();
+
+ private:
+  CArrowStreamTableIterator(ArrowArrayStream* stream, ArrowSchema* schema,
+                      PyObject* context, bool number_to_decimal,
+                      bool force_microsecond_precision);
+
+  /**
+   * Pull the next non-empty batch from the stream into m_currentArray.
+   * Releases the GIL during stream->get_next() so concurrent chunk
+   * downloads (e.g. from S3) are not blocked by the Python interpreter.
+   * Also initialises m_currentArrayView for the conversion step.
+   *
+   * @return true if a batch was loaded; false if exhausted or error
+   */
+  bool loadNextBatch();
+
+  /**
+   * Apply Snowflake type conversions to the current batch via
+   * m_converter, then move the result into the export slots
+   * (m_exportArray / m_exportSchema).
+   * Releases the GIL during the per-column conversion loop, since it is
+   * pure nanoarrow buffer work; any Python exception raised inside the
+   * loop is signalled through py::setPyError(), which re-acquires the GIL
+   * only for the duration of the call. Errors are checked once after the
+   * loop when the GIL is held again.
+   * Sets a Python exception on failure (check py::checkPyError()).
+   */
+  void convertBatch();
+
+  /** Owned handle to the ArrowArrayStream; released on destruction. */
+  std::unique_ptr<ArrowArrayStream, void (*)(ArrowArrayStream*)> m_stream;
+
+  /** Stream schema, read once at construction; immutable across batches. */
+  nanoarrow::UniqueSchema m_schema;
+
+  /** Raw batch from the most recent stream->get_next() call. */
+  nanoarrow::UniqueArray m_currentArray;
+
+  /** View over m_currentArray; rebuilt each loadNextBatch() for the
+   *  converter to read column data through. */
+  nanoarrow::UniqueArrayView m_currentArrayView;
+
+  /** Converted batch ready for PyArrow import.  Populated by
+   *  convertBatch(); consumed by _import_from_c() via getArrowArrayPtr(). */
+  nanoarrow::UniqueArray m_exportArray;
+
+  /** Deep-copied and possibly rewritten schema for the current batch.
+   *  Each batch gets its own copy because conversion may alter child
+   *  schemas (e.g. struct → timestamp).  Consumed via getArrowSchemaPtr(). */
+  nanoarrow::UniqueSchema m_exportSchema;
+
+  /** Deep copy of m_schema used by getStreamSchemaPtr().  Separate from
+   *  m_exportSchema so it can be produced independently of batch iteration. */
+  nanoarrow::UniqueSchema m_streamSchemaExport;
+
+  /** Set to true once stream->get_next() returns an empty release sentinel. */
+  bool m_streamExhausted = false;
+
+  /** Number of top-level columns, cached from schema.n_children. */
+  int64_t m_columnCount = 0;
+
+  /** Python ArrowConverterContext; borrowed reference kept alive by the
+   *  Cython wrapper which stores it as `self.arrow_context`. */
+  PyObject* m_context;
+
+  /** Applies Snowflake-specific type conversions per column. */
+  std::unique_ptr<ArrowTableConverter> m_converter;
+
+  /** Holds the most recent Python exception so the pointer remains valid
+   *  until the next call into the iterator. */
+  py::UniqueRef m_currentPyException;
+
+  static Logger* logger;
+};
+
+}  // namespace sf
+
+#endif  // PC_CARROWSTREAMTABLEITERATOR_HPP