SNOW-86947: update date, time and timestamp converter for arrow

Author: sfc-gh-stakeda

arrow_context.py

@@ -0,0 +1,48 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2012-2019 Snowflake Computing Inc. All right reserved.
+#
+
+from datetime import datetime
+from logging import getLogger
+from .constants import (
+        PARAMETER_TIMEZONE
+)
+
+import pytz
+
+try:
+    import tzlocal
+except ImportError:
+    tzlocal = None
+
+logger = getLogger(__name__)
+
+
+class ArrowConverterContext(object):
+    def __init__(self, session_parameters={}):
+        self._timezone = None if PARAMETER_TIMEZONE not in session_parameters else session_parameters[PARAMETER_TIMEZONE]
+
+    @property
+    def timezone(self):
+        return self._timezone
+
+    @timezone.setter
+    def timezone(self, tz):
+        self._timezone = tz
+
+    def get_session_tz(self):
+        """ Get the session timezone or use the local computer's timezone. """
+        try:
+            tz = 'UTC' if not self.timezone else self.timezone
+            return pytz.timezone(tz)
+        except pytz.exceptions.UnknownTimeZoneError:
+            logger.warning('converting to tzinfo failed')
+            if tzlocal is not None:
+                return tzlocal.get_localzone()
+            else:
+                try:
+                    return datetime.timezone.utc
+                except AttributeError:
+                    return pytz.timezone('UTC')

arrow_iterator.pyx

@@ -8,7 +8,7 @@ from cpython.ref cimport PyObject
 
 cdef extern from "cpp/ArrowIterator/CArrowChunkIterator.hpp" namespace "sf":
     cdef cppclass CArrowChunkIterator:
-        CArrowChunkIterator()
+        CArrowChunkIterator(PyObject* context)
 
         void addRecordBatch(PyObject * rb)
 
@@ -20,16 +20,16 @@ cdef extern from "cpp/ArrowIterator/CArrowChunkIterator.hpp" namespace "sf":
 cdef class PyArrowChunkIterator:
     cdef CArrowChunkIterator * cIterator
 
-    def __cinit__(self, arrow_stream_reader):
-        self.cIterator = new CArrowChunkIterator()
+    def __cinit__(PyArrowChunkIterator self, object arrow_stream_reader, object arrow_context):
+        self.cIterator = new CArrowChunkIterator(<PyObject*>arrow_context)
         for rb in arrow_stream_reader:
             self.cIterator.addRecordBatch(<PyObject *>rb)
         self.cIterator.reset()
 
-    def __dealloc(self):
+    def __dealloc__(PyArrowChunkIterator self):
         del self.cIterator
 
-    def __next__(self):
+    def __next__(PyArrowChunkIterator self):
         ret = <object>self.cIterator.nextRow()
         if ret is None:
             raise StopIteration

arrow_result.pyx

@@ -11,6 +11,7 @@ from .time_util import get_time_millis
 try:
     from pyarrow.ipc import open_stream
     from .arrow_iterator import PyArrowChunkIterator
+    from .arrow_context import ArrowConverterContext
 except ImportError:
     pass
 
@@ -30,6 +31,7 @@ cdef class ArrowResult:
         object _column_idx_to_name
         object _current_chunk_row
         object _chunk_downloader
+        object _arrow_context
 
     def __init__(self, raw_response, cursor):
         self._reset()
@@ -45,7 +47,8 @@ cdef class ArrowResult:
         # result as arrow chunk
         arrow_bytes = b64decode(data.get(u'rowsetBase64'))
         arrow_reader = open_stream(arrow_bytes)
-        self._current_chunk_row = PyArrowChunkIterator(arrow_reader)
+        self._arrow_context = ArrowConverterContext(self._connection._session_parameters)
+        self._current_chunk_row = PyArrowChunkIterator(arrow_reader, self._arrow_context)
 
         if u'chunks' in data:
             chunks = data[u'chunks']
@@ -137,4 +140,5 @@ cdef class ArrowResult:
 
         self._chunk_count = 0
         self._chunk_downloader = None
+        self._arrow_context = None
 

chunk_downloader.py

@@ -21,6 +21,7 @@
 try:
     from pyarrow.ipc import open_stream
     from .arrow_iterator import PyArrowChunkIterator
+    from .arrow_context import ArrowConverterContext
 except ImportError:
     pass
 
@@ -254,7 +255,7 @@ def _fetch_chunk(self, url, headers):
         handler = JsonBinaryHandler(is_raw_binary_iterator=True,
                                     use_ijson=self._use_ijson) \
             if self._query_result_format == 'json' else \
-            ArrowBinaryHandler(self._cursor.description)
+            ArrowBinaryHandler(self._cursor.description, self._connection)
 
         return self._connection.rest.fetch(
             u'get', url, headers,
@@ -321,14 +322,15 @@ def to_iterator(self, raw_data_fd, download_time):
 
 class ArrowBinaryHandler(RawBinaryDataHandler):
 
-    def __init__(self, meta):
+    def __init__(self, meta, connection):
         self._meta = meta
+        self._arrow_context = ArrowConverterContext(connection._session_parameters)
 
     """
     Handler to consume data as arrow stream
     """
     def to_iterator(self, raw_data_fd, download_time):
         gzip_decoder = GzipFile(fileobj=raw_data_fd, mode='r')
         reader = open_stream(gzip_decoder)
-        it = PyArrowChunkIterator(reader)
+        it = PyArrowChunkIterator(reader, self._arrow_context)
         return it

cpp/ArrowIterator/CArrowChunkIterator.cpp

@@ -10,14 +10,16 @@
 #include "BinaryConverter.hpp"
 #include "BooleanConverter.hpp"
 #include "DateConverter.hpp"
+#include "TimeConverter.hpp"
 #include <iostream>
 #include "logging.hpp"
 
 namespace sf
 {
 Logger CArrowChunkIterator::logger("snowflake.connector.CArrowChunkIterator");
 
-CArrowChunkIterator::CArrowChunkIterator() : m_latestReturnedRow(nullptr)
+CArrowChunkIterator::CArrowChunkIterator(PyObject* context)
+: m_latestReturnedRow(nullptr), m_context(context)
 {
 }
 
@@ -232,6 +234,41 @@ void CArrowChunkIterator::initColumnConverters()
         break;
       }
 
+      case SnowflakeType::Type::TIME:
+      {
+        int scale = metaData
+                        ? std::stoi(metaData->value(metaData->FindKey("scale")))
+                        : 9;
+        switch (dt->id())
+        {
+          case arrow::Type::type::INT32:
+          {
+            m_currentBatchConverters.push_back(
+                std::make_shared<sf::TimeConverter<arrow::Int32Array>>(
+                    columnArray, scale));
+            break;
+          }
+
+          case arrow::Type::type::INT64:
+          {
+            m_currentBatchConverters.push_back(
+                std::make_shared<sf::TimeConverter<arrow::Int64Array>>(
+                    columnArray, scale));
+            break;
+          }
+
+          default:
+          {
+            /** cout is playing a placeholder here and will be replaced by
+             * exception soon */
+            std::cout << "unknown arrow internal data type (" << dt->id()
+                      << ") for TIME data" << std::endl;
+            break;
+          }
+        }
+        break;
+      }
+
       default:
       {
         /** cout is playing a placeholder here and will be replaced by exception

cpp/ArrowIterator/CArrowChunkIterator.hpp

@@ -26,7 +26,7 @@ class CArrowChunkIterator
   /**
    * Constructor
    */
-  CArrowChunkIterator();
+  CArrowChunkIterator(PyObject* context);
 
   /**
    * Desctructor
@@ -74,6 +74,9 @@ class CArrowChunkIterator
   /** list of column converters*/
   std::vector<std::shared_ptr<sf::IColumnConverter>> m_currentBatchConverters;
 
+  /** arrow format convert context for the current session */
+  PyObject* m_context;
+
   /**
    * @return python object of tuple which is tuple of all row values
    */

cpp/ArrowIterator/IntConverter.hpp

@@ -5,6 +5,7 @@
 #define PC_INTCONVERTER_HPP
 
 #include "IColumnConverter.hpp"
+#include <iostream>
 
 namespace sf
 {
@@ -18,6 +19,16 @@ class IntConverter : public IColumnConverter
   {
   }
 
+  PyObject* pyLongForward(int64_t value)
+  {
+    return PyLong_FromLongLong(value);
+  }
+
+  PyObject* pyLongForward(int32_t value)
+  {
+    return PyLong_FromLong(value);
+  }
+
   PyObject* toPyObject(int64_t rowIndex) override;
 
 private:
@@ -29,7 +40,8 @@ PyObject* IntConverter<T>::toPyObject(int64_t rowIndex)
 {
   if (m_array->IsValid(rowIndex))
   {
-    return PyLong_FromLongLong(m_array->Value(rowIndex));
+    // TODO : this forward function need to be tested in Win64
+    return pyLongForward(m_array->Value(rowIndex));
   }
   else
   {

cpp/ArrowIterator/TimeConverter.cpp

@@ -0,0 +1,12 @@
+/*
+ * Copyright (c) 2013-2019 Snowflake Computing
+ */
+#include "TimeConverter.hpp"
+#include "Python/Helpers.hpp"
+#include "Util/time.hpp"
+
+namespace sf
+{
+/** this file is here for future use and if this is useless at the end, it will
+ * be removed */
+}  // namespace sf

cpp/ArrowIterator/TimeConverter.hpp

@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) 2013-2019 Snowflake Computing
+ */
+#ifndef PC_TIMECONVERTER_HPP
+#define PC_TIMECONVERTER_HPP
+
+#include "IColumnConverter.hpp"
+#include "Python/Common.hpp"
+#include "Python/Helpers.hpp"
+#include "Util/time.hpp"
+
+namespace sf
+{
+
+template <typename T>
+class TimeConverter : public IColumnConverter
+{
+public:
+  TimeConverter(std::shared_ptr<arrow::Array> array, int32_t scale)
+  : m_array(std::dynamic_pointer_cast<T>(array)), m_scale(scale)
+  {
+  }
+
+  PyObject* toPyObject(int64_t rowIndex) override;
+
+private:
+  /** can be arrow::Int32Array and arrow::Int64Array */
+  std::shared_ptr<T> m_array;
+
+  int32_t m_scale;
+
+  static py::UniqueRef& m_pyDatetimeTime();
+};
+
+template <typename T>
+PyObject* TimeConverter<T>::toPyObject(int64_t rowIndex)
+{
+  if (m_array->IsValid(rowIndex))
+  {
+    int64_t seconds = m_array->Value(rowIndex);
+    using namespace internal;
+    py::PyUniqueLock lock;
+    return PyObject_CallFunction(m_pyDatetimeTime().get(), "iiii",
+                                 getHourFromSeconds(seconds, m_scale),
+                                 getMinuteFromSeconds(seconds, m_scale),
+                                 getSecondFromSeconds(seconds, m_scale),
+                                 getMicrosecondFromSeconds(seconds, m_scale));
+  }
+  else
+  {
+    Py_RETURN_NONE;
+  }
+}
+
+template <typename T>
+py::UniqueRef& TimeConverter<T>::m_pyDatetimeTime()
+{
+  static py::UniqueRef pyDatetimeTime;
+  if (pyDatetimeTime.empty())
+  {
+    py::PyUniqueLock lock;
+    py::UniqueRef pyDatetimeModule;
+    py::importPythonModule("datetime", pyDatetimeModule);
+    /** TODO : to check status here */
+
+    py::importFromModule(pyDatetimeModule, "time", pyDatetimeTime);
+  }
+  return pyDatetimeTime;
+}
+
+}  // namespace sf
+
+#endif  // PC_TIMECONVERTER_HPP