SNOW-87349: Use cython to generate glue code from python to c++ (not enabled)

Author: sfc-gh-stakeda

arrow_iterator.pyx

@@ -0,0 +1,28 @@
+# distutils: language = c++
+
+from cpython.ref cimport PyObject
+
+cdef extern from "cpp/ArrowIterator/CArrowChunkIterator.hpp" namespace "sf":
+    cdef cppclass CArrowChunkIterator:
+        CArrowChunkIterator()
+
+        void addRecordBatch(PyObject * rb)
+
+        PyObject *nextRow();
+
+
+cdef class PyArrowChunkIterator:
+    cdef CArrowChunkIterator thisptr
+
+    def __cinit__(self):
+        self.thisptr = CArrowChunkIterator()
+
+    def add_record_batch(self, rb):
+        self.thisptr.addRecordBatch(<PyObject *>rb)
+
+    def __next__(self):
+        ret = <object>self.thisptr.nextRow()
+        if ret is None:
+            raise StopIteration
+        else:
+            return ret

arrow_result.pyx

@@ -10,7 +10,7 @@ from .telemetry import TelemetryField
 from .time_util import get_time_millis
 try:
     from pyarrow.ipc import open_stream
-    from .libarrow_iterator import ArrowChunkIterator
+    from .arrow_iterator import PyArrowChunkIterator
 except ImportError:
     pass
 
@@ -45,7 +45,7 @@ cdef class ArrowResult:
         # result as arrow chunk
         arrow_bytes = b64decode(data.get(u'rowsetBase64'))
         arrow_reader = open_stream(arrow_bytes)
-        self._current_chunk_row = ArrowChunkIterator()
+        self._current_chunk_row = PyArrowChunkIterator()
         for rb in arrow_reader:
             self._current_chunk_row.add_record_batch(rb)
 

chunk_downloader.py

@@ -20,7 +20,7 @@
 
 try:
     from pyarrow.ipc import open_stream
-    from .arrow_iterator import ArrowChunkIterator
+    from .arrow_iterator import PyArrowChunkIterator
 except ImportError:
     pass
 
@@ -345,7 +345,7 @@ def __init__(self, meta):
     def to_iterator(self, raw_data_fd, download_time):
         gzip_decoder = GzipFile(fileobj=raw_data_fd, mode='r')
         reader = open_stream(gzip_decoder)
-        it = ArrowChunkIterator()
+        it = PyArrowChunkIterator()
         for rb in reader:
             it.add_record_batch(rb)
         return it

cpp/ArrowIterator/CArrowChunkIterator.cpp

@@ -10,9 +10,12 @@ sf::CArrowChunkIterator::CArrowChunkIterator()
     this->reset();
 }
 
-void sf::CArrowChunkIterator::addRecordBatch(std::shared_ptr<arrow::RecordBatch> rb)
+void sf::CArrowChunkIterator::addRecordBatch(PyObject * rb)
 {
-    m_cRecordBatches.push_back(rb);
+    std::shared_ptr<arrow::RecordBatch> cRecordBatch;
+    arrow::Status status = arrow::py::unwrap_record_batch(rb, &cRecordBatch);
+
+    m_cRecordBatches.push_back(cRecordBatch);
     m_columnCount = m_cRecordBatches[0]->num_columns();
     m_batchCount = m_cRecordBatches.size();
 }
@@ -44,7 +47,7 @@ PyObject * sf::CArrowChunkIterator::nextRow()
         }
     }
 
-    return NULL;
+    return Py_None;
 }
 
 PyObject * sf::CArrowChunkIterator::currentRowAsTuple()
@@ -68,6 +71,22 @@ void sf::CArrowChunkIterator::initColumnConverters()
         std::shared_ptr<arrow::DataType> dt = schema->field(i)->type();
         switch(dt->id())
         {
+
+            case arrow::Type::type::INT8:
+                m_currentBatchConverters.push_back(
+                        std::make_shared<sf::Int8Converter>(columnArray.get()));
+                break;
+
+            case arrow::Type::type::INT16:
+                m_currentBatchConverters.push_back(
+                        std::make_shared<sf::Int16Converter>(columnArray.get()));
+                break;
+
+            case arrow::Type::type::INT32:
+                m_currentBatchConverters.push_back(
+                        std::make_shared<sf::Int32Converter>(columnArray.get()));
+                break;
+
             case arrow::Type::type::INT64:
                 m_currentBatchConverters.push_back(
                     std::make_shared<sf::Int64Converter>(columnArray.get()));

cpp/ArrowIterator/CArrowChunkIterator.hpp

@@ -34,7 +34,7 @@ class CArrowChunkIterator
      * Add Arrow RecordBach to current chunk
      * @param rb recordbatch to be added
      */
-    void addRecordBatch(std::shared_ptr<arrow::RecordBatch> rb);
+    void addRecordBatch(PyObject *rb);
 
     /**
      * @return a python tuple object which contains all data in current row

cpp/ArrowIterator/IntConverter.cpp

@@ -13,3 +13,35 @@ PyObject* sf::Int64Converter::toPyObject(int64_t rowIndex)
     return (m_array->IsValid(rowIndex)) ? PyLong_FromLong(m_array->Value(rowIndex)) : Py_None;
 }
 
+sf::Int32Converter::Int32Converter(arrow::Array *array)
+{
+    m_array = dynamic_cast<arrow::Int32Array *>(array);
+}
+
+PyObject* sf::Int32Converter::toPyObject(int64_t rowIndex)
+{
+    return (m_array->IsValid(rowIndex)) ? PyLong_FromLong(m_array->Value(rowIndex)) : Py_None;
+}
+
+
+sf::Int16Converter::Int16Converter(arrow::Array *array)
+{
+    m_array = dynamic_cast<arrow::Int16Array *>(array);
+}
+
+PyObject* sf::Int16Converter::toPyObject(int64_t rowIndex)
+{
+    return (m_array->IsValid(rowIndex)) ? PyLong_FromLong(m_array->Value(rowIndex)) : Py_None;
+}
+
+sf::Int8Converter::Int8Converter(arrow::Array *array)
+{
+    m_array = dynamic_cast<arrow::Int8Array *>(array);
+}
+
+PyObject* sf::Int8Converter::toPyObject(int64_t rowIndex)
+{
+    return (m_array->IsValid(rowIndex)) ? PyLong_FromLong(m_array->Value(rowIndex)) : Py_None;
+}
+
+

cpp/ArrowIterator/IntConverter.hpp

@@ -15,12 +15,45 @@ class Int64Converter : public IColumnConverter
 public:
     Int64Converter(arrow::Array * array);
 
-    PyObject * toPyObject(long rowIndex) override;
+    PyObject * toPyObject(int64_t rowIndex) override;
 
 private:
     arrow::Int64Array * m_array;
 };
 
+class Int32Converter : public IColumnConverter
+{
+public:
+    Int32Converter(arrow::Array * array);
+
+    PyObject * toPyObject(int64_t rowIndex) override;
+
+private:
+    arrow::Int32Array * m_array;
+};
+
+class Int16Converter : public IColumnConverter
+{
+public:
+    Int16Converter(arrow::Array * array);
+
+    PyObject * toPyObject(int64_t rowIndex) override;
+
+private:
+    arrow::Int16Array * m_array;
+};
+
+class Int8Converter : public IColumnConverter
+{
+public:
+    Int8Converter(arrow::Array * array);
+
+    PyObject * toPyObject(int64_t rowIndex) override;
+
+private:
+    arrow::Int8Array * m_array;
+};
+
 }
 
 #endif

cpp/ArrowIterator/PyArrowChunkIteratorDef.cpp

@@ -14,7 +14,7 @@ class StringConverter : public IColumnConverter
 public:
     StringConverter(arrow::Array * array);
 
-    PyObject * toPyObject(long rowIndex) override;
+    PyObject * toPyObject(int64_t rowIndex) override;
 
 private:
     arrow::StringArray * m_array;