SNOW-87349: Outupt cython code-generated .cpp code to build directory

Author: sfc-gh-stakeda

arrow_iterator.pyx

@@ -1,3 +1,7 @@
+#
+# Copyright (c) 2012-2019 Snowflake Computing Inc. All right reserved.
+#
+
 # distutils: language = c++
 
 from cpython.ref cimport PyObject

cpp/ArrowIterator/CArrowChunkIterator.cpp

@@ -4,8 +4,9 @@
 #include "CArrowChunkIterator.hpp"
 #include "IntConverter.hpp"
 #include "StringConverter.hpp"
+#include "FloatConverter.hpp"
 
-sf::CArrowChunkIterator::CArrowChunkIterator()
+sf::CArrowChunkIterator::CArrowChunkIterator() : m_latestReturnedRow(nullptr)
 {
     this->reset();
 }
@@ -25,11 +26,15 @@ void sf::CArrowChunkIterator::reset()
     m_currentBatchIndex = -1;
     m_rowIndexInBatch = -1;
     m_rowCountInBatch = 0;
+    Py_XDECREF(m_latestReturnedRow);
+    m_latestReturnedRow = nullptr;
 }
 
 PyObject * sf::CArrowChunkIterator::nextRow()
 {
     m_rowIndexInBatch ++;
+    Py_XDECREF(m_latestReturnedRow);
+    m_latestReturnedRow = nullptr;
 
     if (m_rowIndexInBatch < m_rowCountInBatch)
     {
@@ -57,7 +62,7 @@ PyObject * sf::CArrowChunkIterator::currentRowAsTuple()
     {
         PyTuple_SET_ITEM(tuple, i, m_currentBatchConverters[i]->toPyObject(m_rowIndexInBatch));
     }
-    return tuple;
+    return m_latestReturnedRow = tuple;
 }
 
 void sf::CArrowChunkIterator::initColumnConverters()
@@ -97,6 +102,11 @@ void sf::CArrowChunkIterator::initColumnConverters()
                     std::make_shared<sf::StringConverter>(columnArray.get()));
                 break;
 
+            case arrow::Type::type::DOUBLE:
+                m_currentBatchConverters.push_back(
+                    std::make_shared<sf::FloatConverter>(columnArray.get()));
+                break;
+
             default:
                 break;
         }

cpp/ArrowIterator/CArrowChunkIterator.hpp

@@ -28,7 +28,10 @@ class CArrowChunkIterator
     /**
      * Desctructor
      */
-    ~CArrowChunkIterator() = default;
+    ~CArrowChunkIterator()
+    {
+        Py_XDECREF(m_latestReturnedRow);
+    }
 
     /**
      * Add Arrow RecordBach to current chunk
@@ -59,6 +62,9 @@ class CArrowChunkIterator
 
     /** total number of rows inside current record batch */
     int64_t m_rowCountInBatch;
+   
+    /** pointer to the latest returned python tuple(row) result */
+    PyObject* m_latestReturnedRow;
 
     /** list of column converters*/
     std::vector<std::shared_ptr<sf::IColumnConverter>> m_currentBatchConverters;

cpp/ArrowIterator/FloatConverter.cpp

@@ -0,0 +1,16 @@
+/*
+ * Copyright (c) 2013-2019 Snowflake Computing
+ */
+#include "FloatConverter.hpp"
+
+sf::FloatConverter::FloatConverter(arrow::Array* array)
+{
+    /** snowflake float is 64-precision, which refers to double here */
+    m_array = dynamic_cast<arrow::DoubleArray*>(array);
+}
+
+PyObject* sf::FloatConverter::toPyObject(int64_t rowIndex)
+{
+    return (m_array->IsValid(rowIndex)) ? PyFloat_FromDouble(m_array->Value(rowIndex)) : Py_None;
+}
+

cpp/ArrowIterator/FloatConverter.hpp

@@ -0,0 +1,24 @@
+/*
+ * Copyright (c) 2013-2019 Snowflake Computing
+ */
+#ifndef PC_FLOATCONVERTER_HPP
+#define PC_FLOATCONVERTER_HPP
+
+#include "IColumnConverter.hpp"
+
+namespace sf
+{
+
+class FloatConverter : public IColumnConverter
+{
+public:
+    explicit FloatConverter(arrow::Array* array);
+
+    PyObject* toPyObject(long rowIndex) override;
+
+private:
+    arrow::DoubleArray* m_array;
+};
+}
+
+#endif

cpp/ArrowIterator/IntConverter.hpp

@@ -6,7 +6,6 @@
 
 #include "IColumnConverter.hpp"
 
-
 namespace sf
 {
 

setup.py

@@ -7,6 +7,7 @@
 from os import path
 import os
 import sys
+from sys import platform
 from shutil import copy
 import glob
 
@@ -38,20 +39,16 @@
 
 if isBuildExtEnabled == 'true':
     from Cython.Distutils import build_ext
+    from Cython.Build import cythonize
     import os
     import pyarrow
 
-    extensions = [
-        Extension(name='snowflake.connector.arrow_result', sources=['arrow_result.pyx']),
-        Extension(name='snowflake.connector.arrow_iterator', sources=['arrow_iterator.pyx',
-                                                                      'cpp/ArrowIterator/CArrowChunkIterator.cpp',
-                                                                      'cpp/ArrowIterator/IntConverter.cpp',
-                                                                      'cpp/ArrowIterator/StringConverter.cpp'],
-                  extra_compile_args=['-std=c++11'],
-                  extra_link_args=['-Wl,-rpath,$ORIGIN'],
-                  language='c++'
-                  ),
-    ]
+    extensions = cythonize(
+        [
+            Extension(name='snowflake.connector.arrow_iterator', sources=['arrow_iterator.pyx']),
+            Extension(name='snowflake.connector.arrow_result', sources=['arrow_result.pyx'])
+        ],
+        build_dir=os.path.join('build', 'cython'))
 
     class MyBuildExt(build_ext):
 
@@ -61,36 +58,50 @@ def build_extension(self, ext):
             if ext.name == 'snowflake.connector.arrow_iterator':
                 self._copy_arrow_lib()
 
-                ext.include_dirs.append(self._get_arrow_include_dir())
+                ext.sources += ['cpp/ArrowIterator/CArrowChunkIterator.cpp',
+                                'cpp/ArrowIterator/FloatConverter.cpp',
+                                'cpp/ArrowIterator/IntConverter.cpp',
+                                'cpp/ArrowIterator/StringConverter.cpp']
+                ext.include_dirs.append('cpp/ArrowIterator/')
+                ext.include_dirs.append(pyarrow.get_include())
+
+                ext.extra_compile_args.append('-std=c++11')
+
                 ext.library_dirs.append(os.path.join(current_dir, self.build_lib, 'snowflake', 'connector'))
                 ext.extra_link_args += self._get_arrow_lib_as_linker_input()
+                ext.extra_link_args += ['-Wl,-rpath,$ORIGIN']
 
             build_ext.build_extension(self, ext)
 
-        def _get_arrow_include_dir(self):
-            return pyarrow.get_include()
-
         def _get_arrow_lib_dir(self):
             return pyarrow.get_library_dirs()[0]
 
         def _copy_arrow_lib(self):
             arrow_lib = pyarrow.get_libraries() + \
                         ['arrow_flight', 'arrow_boost_regex', 'arrow_boost_system', 'arrow_boost_filesystem']
             for lib in arrow_lib:
-                lib_pattern = '{}/lib{}.so*'.format(self._get_arrow_lib_dir(), lib)
+                lib_pattern = self._get_pyarrow_lib_pattern(lib)
                 source = glob.glob(lib_pattern)[0]
                 copy(source, os.path.join(self.build_lib, 'snowflake', 'connector'))
 
         def _get_arrow_lib_as_linker_input(self):
             arrow_lib = pyarrow.get_libraries()
             link_lib = []
             for lib in arrow_lib:
-                lib_pattern = '{}/lib{}.so*'.format(self._get_arrow_lib_dir(), lib)
+                lib_pattern = self._get_pyarrow_lib_pattern(lib)
                 source = glob.glob(lib_pattern)[0]
                 link_lib.append(source)
 
             return link_lib
 
+        def _get_pyarrow_lib_pattern(self, lib_name):
+            if platform.startswith('linux'):
+                return '{}/lib{}.so*'.format(self._get_arrow_lib_dir(), lib_name)
+            elif platform == 'darwin':
+                return '{}/lib{}*dylib'.format(self._get_arrow_lib_dir(), lib_name)
+            else:
+                raise RuntimeError('Building on platform {} is not supported yet.'.format(platform))
+
     cmd_class = {
         "build_ext": MyBuildExt
     }
@@ -165,8 +176,8 @@ def _get_arrow_lib_as_linker_input(self):
             'keyring!=16.1.0'
         ],
         "arrow-result": [
-            'pyarrow>=0.13.0;python_version>"3.4"',
-            'pyarrow>=0.13.0;python_version<"3.0"'
+            'pyarrow>=0.14.0;python_version>"3.4"',
+            'pyarrow>=0.14.0;python_version<"3.0"'
         ]
     },
 

test/test_unit_arrow_chunk_iterator.py

@@ -124,3 +124,56 @@ def test_iterate_over_int64_chunk():
             assert count == 100
             break
 
+
+@pytest.mark.skip(
+    reason="Cython is not enabled in build env")
+def test_iterate_over_float_chunk():
+    stream = BytesIO()
+    field_foo = pyarrow.field("column_foo", pyarrow.float64(), True)
+    field_bar = pyarrow.field("column_bar", pyarrow.float64(), True)
+    schema = pyarrow.schema([field_foo, field_bar])
+    column_meta = [
+      ("column_foo", "FLOAT", None, 0, 0, 0, 0),
+      ("column_bar", "FLOAT", None, 0, 0, 0, 0)
+    ]
+
+    column_size = 2
+    batch_row_count = 10
+    batch_count = 10
+    expected_data = []
+    writer = RecordBatchStreamWriter(stream, schema)
+
+    for i in range(batch_count):
+        column_arrays = []
+        py_arrays = []
+        for j in range(column_size):
+            column_data = []
+            for k in range(batch_row_count):
+                data = None if bool(random.getrandbits(1)) else random.uniform(-100.0, 100.0)
+                column_data.append(data)
+            column_arrays.append(column_data)
+            py_arrays.append(pyarrow.array(column_data))
+
+        expected_data.append(column_arrays)
+        rb = RecordBatch.from_arrays(py_arrays, ["column_foo", "column_bar"])
+        writer.write_batch(rb)
+
+    writer.close()
+
+    # seek stream to begnning so that we can read from stream
+    stream.seek(0)
+    reader = RecordBatchStreamReader(stream)
+    it = PyArrowChunkIterator()
+    for rb in reader:
+        it.add_record_batch(rb)
+
+    count = 0
+    while True:
+        try:
+            val = next(it)
+            assert val[0] == expected_data[int(count / 10)][0][count % 10]
+            assert val[1] == expected_data[int(count / 10)][1][count % 10]
+            count += 1
+        except StopIteration:
+            assert count == 100
+            break