SNOW-86947: added timestamp converter for arrow

Author: sfc-gh-stakeda

arrow_context.py

@@ -4,11 +4,13 @@
 # Copyright (c) 2012-2019 Snowflake Computing Inc. All right reserved.
 #
 
-from datetime import datetime
+import time
+from datetime import datetime, timedelta
 from logging import getLogger
 from .constants import (
-        PARAMETER_TIMEZONE
-)
+    PARAMETER_TIMEZONE)
+from .converter import (
+    _generate_tzinfo_from_tzoffset)
 
 import pytz
 
@@ -17,6 +19,8 @@
 except ImportError:
     tzlocal = None
 
+ZERO_EPOCH = datetime.utcfromtimestamp(0)
+
 logger = getLogger(__name__)
 
 
@@ -32,7 +36,7 @@ def timezone(self):
     def timezone(self, tz):
         self._timezone = tz
 
-    def get_session_tz(self):
+    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
@@ -46,3 +50,46 @@ def get_session_tz(self):
                     return datetime.timezone.utc
                 except AttributeError:
                     return pytz.timezone('UTC')
+
+    def TIMESTAMP_TZ_to_python(self, microseconds, tz):
+        """
+        TIMESTAMP TZ to datetime
+
+        The timezone offset is piggybacked
+
+        @para microseconds : float
+        @para tz : int
+        """
+
+        tzinfo = _generate_tzinfo_from_tzoffset(tz - 1440)
+        return datetime.fromtimestamp(microseconds, tz=tzinfo)
+
+    def TIMESTAMP_TZ_to_python_windows(self, microseconds, tz):
+        tzinfo = _generate_tzinfo_from_tzoffset(tz - 1440)
+        t = ZERO_EPOCH + timedelta(seconds=microseconds)
+        if pytz.utc != tzinfo:
+            t += tzinfo.utcoffset(t, is_dst=False)
+        return t.replace(tzinfo=tzinfo)
+
+    def TIMESTAMP_NTZ_to_python(self, microseconds):
+        return datetime.utcfromtimestamp(microseconds)
+
+    def TIMESTAMP_NTZ_to_python_windows(self, microseconds):
+        return ZERO_EPOCH + timedelta(seconds=(microseconds))
+
+    def TIMESTAMP_LTZ_to_python(self, microseconds):
+        tzinfo = self._get_session_tz()
+        return datetime.fromtimestamp(microseconds, tz=tzinfo)
+
+    def TIMESTAMP_LTZ_to_python_windows(self, microseconds):
+        tzinfo = self._get_session_tz()
+        try:
+            t0 = ZERO_EPOCH + timedelta(seconds=(microseconds))
+            t = pytz.utc.localize(t0, is_dst=False).astimezone(tzinfo)
+            return t
+        except OverflowError:
+            logger.debug(
+                "OverflowError in converting from epoch time to "
+                "timestamp_ltz: %s(ms). Falling back to use struct_time."
+            )
+            return time.localtime(microseconds)

cpp/ArrowIterator/BinaryConverter.cpp

@@ -5,12 +5,14 @@
 
 namespace sf
 {
+Logger BinaryConverter::logger("snowflake.connector.BinaryConverter");
+
 BinaryConverter::BinaryConverter(std::shared_ptr<arrow::Array> array)
 : m_array(std::dynamic_pointer_cast<arrow::BinaryArray>(array))
 {
 }
 
-PyObject* BinaryConverter::toPyObject(int64_t rowIndex)
+PyObject* BinaryConverter::toPyObject(int64_t rowIndex) const
 {
   if (m_array->IsValid(rowIndex))
   {

cpp/ArrowIterator/BinaryConverter.hpp

@@ -5,19 +5,22 @@
 #define PC_BINARYCONVERTER_HPP
 
 #include "IColumnConverter.hpp"
+#include "logging.hpp"
 
 namespace sf
 {
 
 class BinaryConverter : public IColumnConverter
 {
 public:
-  BinaryConverter(std::shared_ptr<arrow::Array> array);
+  explicit BinaryConverter(std::shared_ptr<arrow::Array> array);
 
-  PyObject* toPyObject(int64_t rowIndex) override;
+  PyObject* toPyObject(int64_t rowIndex) const override;
 
 private:
   std::shared_ptr<arrow::BinaryArray> m_array;
+
+  static Logger logger;
 };
 
 }  // namespace sf

cpp/ArrowIterator/BooleanConverter.cpp

@@ -3,8 +3,6 @@
  */
 #include "BooleanConverter.hpp"
 
-#include <iostream>
-
 namespace sf
 {
 
@@ -13,7 +11,7 @@ BooleanConverter::BooleanConverter(std::shared_ptr<arrow::Array> array)
 {
 }
 
-PyObject* BooleanConverter::toPyObject(int64_t rowIndex)
+PyObject* BooleanConverter::toPyObject(int64_t rowIndex) const
 {
   if (m_array->IsValid(rowIndex))
   {

cpp/ArrowIterator/BooleanConverter.hpp

@@ -14,7 +14,7 @@ class BooleanConverter : public IColumnConverter
 public:
   explicit BooleanConverter(std::shared_ptr<arrow::Array> array);
 
-  PyObject* toPyObject(int64_t rowIndex) override;
+  PyObject* toPyObject(int64_t rowIndex) const override;
 
 private:
   std::shared_ptr<arrow::BooleanArray> m_array;

cpp/ArrowIterator/CArrowChunkIterator.cpp

@@ -10,8 +10,8 @@
 #include "BinaryConverter.hpp"
 #include "BooleanConverter.hpp"
 #include "DateConverter.hpp"
+#include "TimeStampConverter.hpp"
 #include "TimeConverter.hpp"
-#include <iostream>
 #include "logging.hpp"
 
 namespace sf
@@ -37,8 +37,7 @@ void CArrowChunkIterator::reset()
   m_currentBatchIndex = -1;
   m_rowIndexInBatch = -1;
   m_rowCountInBatch = 0;
-  Py_XDECREF(m_latestReturnedRow);
-  m_latestReturnedRow = nullptr;
+  m_latestReturnedRow.reset();
 
   logger.info("Arrow chunk info: batchCount %d, columnCount %d", m_batchCount,
               m_columnCount);
@@ -47,12 +46,11 @@ void CArrowChunkIterator::reset()
 PyObject* CArrowChunkIterator::nextRow()
 {
   m_rowIndexInBatch++;
-  Py_XDECREF(m_latestReturnedRow);
-  m_latestReturnedRow = nullptr;
 
   if (m_rowIndexInBatch < m_rowCountInBatch)
   {
-    return this->currentRowAsTuple();
+    this->currentRowAsTuple();
+    return m_latestReturnedRow.get();
   }
   else
   {
@@ -66,7 +64,8 @@ PyObject* CArrowChunkIterator::nextRow()
       logger.info("Current batch index: %d, rows in current batch: %d",
                   m_currentBatchIndex, m_rowCountInBatch);
 
-      return this->currentRowAsTuple();
+      this->currentRowAsTuple();
+      return m_latestReturnedRow.get();
     }
   }
 
@@ -75,15 +74,16 @@ PyObject* CArrowChunkIterator::nextRow()
   return Py_None;
 }
 
-PyObject* CArrowChunkIterator::currentRowAsTuple()
+void CArrowChunkIterator::currentRowAsTuple()
 {
-  PyObject* tuple = PyTuple_New(m_columnCount);
+  m_latestReturnedRow.reset(PyTuple_New(m_columnCount));
   for (int i = 0; i < m_columnCount; i++)
   {
     PyTuple_SET_ITEM(
-        tuple, i, m_currentBatchConverters[i]->toPyObject(m_rowIndexInBatch));
+        m_latestReturnedRow.get(), i,
+        m_currentBatchConverters[i]->toPyObject(m_rowIndexInBatch));
   }
-  return m_latestReturnedRow = tuple;
+  return;
 }
 
 void CArrowChunkIterator::initColumnConverters()
@@ -189,10 +189,9 @@ void CArrowChunkIterator::initColumnConverters()
 
           default:
           {
-            /** cout is playing a placeholder here and will be replaced by
-             * exception soon */
-            std::cout << "unknown arrow internal data type (" << dt->id()
-                      << ") for FIXED data" << std::endl;
+            /** TODO: how to throw an exception will be decided later */
+            logger.error("unknown arrow internal data type(%d) for FIXED data",
+                         dt->id());
             break;
           }
         }
@@ -259,23 +258,130 @@ void CArrowChunkIterator::initColumnConverters()
 
           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;
+            /** TODO: how to throw an exception will be decided later */
+            logger.error("unknown arrow internal data type(%d) for TIME data",
+                         dt->id());
             break;
           }
         }
         break;
       }
 
+      case SnowflakeType::Type::TIMESTAMP_NTZ:
+      {
+        int scale = metaData
+                        ? std::stoi(metaData->value(metaData->FindKey("scale")))
+                        : 9;
+        switch (dt->id())
+        {
+          case arrow::Type::type::INT64:
+          {
+            m_currentBatchConverters.push_back(
+                std::make_shared<sf::OneFieldTimeStampNTZConverter>(
+                    columnArray, scale, m_context));
+            break;
+          }
+
+          case arrow::Type::type::STRUCT:
+          {
+            m_currentBatchConverters.push_back(
+                std::make_shared<sf::TwoFieldTimeStampNTZConverter>(
+                    columnArray, scale, m_context));
+            break;
+          }
+
+          default:
+          {
+            /** TODO: how to throw an exception will be decided later */
+            logger.error(
+                "unknown arrow internal data type(%d) for TIMESTAMP_NTZ data",
+                dt->id());
+            break;
+          }
+        }
+        break;
+      }
+
+      case SnowflakeType::Type::TIMESTAMP_LTZ:
+      {
+        int scale = metaData
+                        ? std::stoi(metaData->value(metaData->FindKey("scale")))
+                        : 9;
+        switch (dt->id())
+        {
+          case arrow::Type::type::INT64:
+          {
+            m_currentBatchConverters.push_back(
+                std::make_shared<sf::OneFieldTimeStampLTZConverter>(
+                    columnArray, scale, m_context));
+            break;
+          }
+
+          case arrow::Type::type::STRUCT:
+          {
+            m_currentBatchConverters.push_back(
+                std::make_shared<sf::TwoFieldTimeStampLTZConverter>(
+                    columnArray, scale, m_context));
+            break;
+          }
+
+          default:
+          {
+            /** TODO: how to throw an exception will be decided later */
+            logger.error(
+                "unknown arrow internal data type(%d) for TIMESTAMP_LTZ data",
+                dt->id());
+            break;
+          }
+        }
+        break;
+      }
+
+      case SnowflakeType::Type::TIMESTAMP_TZ:
+      {
+        int scale = metaData
+                        ? std::stoi(metaData->value(metaData->FindKey("scale")))
+                        : 9;
+        int byteLength =
+            metaData
+                ? std::stoi(metaData->value(metaData->FindKey("byteLength")))
+                : 16;
+        switch (byteLength)
+        {
+          case 8:
+          {
+            m_currentBatchConverters.push_back(
+                std::make_shared<sf::TwoFieldTimeStampTZConverter>(
+                    columnArray, scale, m_context));
+            break;
+          }
+
+          case 16:
+          {
+            m_currentBatchConverters.push_back(
+                std::make_shared<sf::ThreeFieldTimeStampTZConverter>(
+                    columnArray, scale, m_context));
+            break;
+          }
+
+          default:
+          {
+            /** TODO: how to throw an exception will be decided later */
+            logger.error(
+                "unknown arrow internal data type(%d) for TIMESTAMP_TZ data",
+                dt->id());
+            break;
+          }
+        }
+
+        break;
+      }
+
       default:
       {
-        /** cout is playing a placeholder here and will be replaced by exception
-         * soon */
-        std::cout << "[ERROR] unknown snowflake data type : "
-                  << metaData->value(metaData->FindKey("logicalType"))
-                  << std::endl;
+        /** TODO: how to throw an exception will be decided later */
+        logger.error("unknown snowflake data type : %d",
+                     metaData->value(metaData->FindKey("logicalType")));
         break;
       }
     }