SNOW-2052629: Add basic arrow support for Interval data types (#2296)

(cherry picked from commit bbc2f80)

Author: sfc-gh-nkumar

setup.py

@@ -103,6 +103,7 @@ def build_extension(self, ext):
                         "FixedSizeListConverter.cpp",
                         "FloatConverter.cpp",
                         "IntConverter.cpp",
+                        "IntervalConverter.cpp",
                         "MapConverter.cpp",
                         "ObjectConverter.cpp",
                         "SnowflakeType.cpp",

src/snowflake/connector/arrow_context.py

@@ -15,7 +15,7 @@
 from .converter import _generate_tzinfo_from_tzoffset
 
 if TYPE_CHECKING:
-    from numpy import datetime64, float64, int64
+    from numpy import datetime64, float64, int64, timedelta64
 
 
 try:
@@ -163,3 +163,41 @@ def DECFLOAT_to_decimal(self, exponent: int, significand: bytes) -> decimal.Deci
 
     def DECFLOAT_to_numpy_float64(self, exponent: int, significand: bytes) -> float64:
         return numpy.float64(self.DECFLOAT_to_decimal(exponent, significand))
+
+    def INTERVAL_YEAR_MONTH_to_numpy_timedelta(self, months: int) -> timedelta64:
+        return numpy.timedelta64(months, "M")
+
+    def INTERVAL_DAY_TIME_int_to_numpy_timedelta(self, nanos: int) -> timedelta64:
+        return numpy.timedelta64(nanos, "ns")
+
+    def INTERVAL_DAY_TIME_int_to_timedelta(self, nanos: int) -> timedelta:
+        # Python timedelta only supports microsecond precision. We receive value in
+        # nanoseconds.
+        return timedelta(microseconds=nanos // 1000)
+
+    def INTERVAL_DAY_TIME_decimal_to_numpy_timedelta(self, value: bytes) -> timedelta64:
+        # Snowflake supports up to 9 digits leading field precision for the day-time
+        # interval. That when represented in nanoseconds can not be stored in a 64-bit
+        # integer. So we send these as Decimal128 from server to client.
+        # Arrow uses little-endian by default.
+        # https://arrow.apache.org/docs/format/Columnar.html#byte-order-endianness
+        nanos = int.from_bytes(value, byteorder="little", signed=True)
+        # Numpy timedelta only supports up to 64-bit integers, so we need to change the
+        # unit to milliseconds to avoid overflow.
+        # Max value received from server
+        #   = 10**9 * NANOS_PER_DAY - 1
+        #   = 86399999999999999999999 nanoseconds
+        #   = 86399999999999999 milliseconds
+        # math.log2(86399999999999999) = 56.3 < 64
+        return numpy.timedelta64(nanos // 1_000_000, "ms")
+
+    def INTERVAL_DAY_TIME_decimal_to_timedelta(self, value: bytes) -> timedelta:
+        # Snowflake supports up to 9 digits leading field precision for the day-time
+        # interval. That when represented in nanoseconds can not be stored in a 64-bit
+        # integer. So we send these as Decimal128 from server to client.
+        # Arrow uses little-endian by default.
+        # https://arrow.apache.org/docs/format/Columnar.html#byte-order-endianness
+        nanos = int.from_bytes(value, byteorder="little", signed=True)
+        # Python timedelta only supports microsecond precision. We receive value in
+        # nanoseconds.
+        return timedelta(microseconds=nanos // 1000)

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowChunkIterator.cpp

@@ -13,6 +13,7 @@
 #include "FixedSizeListConverter.hpp"
 #include "FloatConverter.hpp"
 #include "IntConverter.hpp"
+#include "IntervalConverter.hpp"
 #include "MapConverter.hpp"
 #include "ObjectConverter.hpp"
 #include "StringConverter.hpp"
@@ -479,6 +480,36 @@ std::shared_ptr<sf::IColumnConverter> getConverterFromSchema(
       break;
     }
 
+    case SnowflakeType::Type::INTERVAL_YEAR_MONTH: {
+      converter = std::make_shared<sf::IntervalYearMonthConverter>(
+          array, context, useNumpy);
+      break;
+    }
+
+    case SnowflakeType::Type::INTERVAL_DAY_TIME: {
+      switch (schemaView.type) {
+        case NANOARROW_TYPE_INT64:
+          converter = std::make_shared<sf::IntervalDayTimeConverterInt>(
+              array, context, useNumpy);
+          break;
+        case NANOARROW_TYPE_DECIMAL128:
+          converter = std::make_shared<sf::IntervalDayTimeConverterDecimal>(
+              array, context, useNumpy);
+          break;
+        default: {
+          std::string errorInfo = Logger::formatString(
+              "[Snowflake Exception] unknown arrow internal data type(%d) "
+              "for OBJECT data in %s",
+              NANOARROW_TYPE_ENUM_STRING[schemaView.type],
+              schemaView.schema->name);
+          logger->error(__FILE__, __func__, __LINE__, errorInfo.c_str());
+          PyErr_SetString(PyExc_Exception, errorInfo.c_str());
+          break;
+        }
+      }
+      break;
+    }
+
     default: {
       std::string errorInfo = Logger::formatString(
           "[Snowflake Exception] unknown snowflake data type : %d", st);

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/IntervalConverter.cpp

@@ -0,0 +1,71 @@
+#include "IntervalConverter.hpp"
+
+#include <memory>
+#include <string>
+
+#include "Python/Common.hpp"
+#include "Python/Helpers.hpp"
+
+namespace sf {
+
+static constexpr char INTERVAL_DT_DECIMAL_TO_NUMPY_TIMEDELTA[] =
+    "INTERVAL_DAY_TIME_decimal_to_numpy_timedelta";
+static constexpr char INTERVAL_DT_DECIMAL_TO_TIMEDELTA[] =
+    "INTERVAL_DAY_TIME_decimal_to_timedelta";
+static constexpr char INTERVAL_DT_INT_TO_NUMPY_TIMEDELTA[] =
+    "INTERVAL_DAY_TIME_int_to_numpy_timedelta";
+static constexpr char INTERVAL_DT_INT_TO_TIMEDELTA[] =
+    "INTERVAL_DAY_TIME_int_to_timedelta";
+
+IntervalYearMonthConverter::IntervalYearMonthConverter(ArrowArrayView* array,
+                                                       PyObject* context,
+                                                       bool useNumpy)
+    : m_array(array), m_context(context), m_useNumpy(useNumpy) {}
+
+PyObject* IntervalYearMonthConverter::toPyObject(int64_t rowIndex) const {
+  if (ArrowArrayViewIsNull(m_array, rowIndex)) {
+    Py_RETURN_NONE;
+  }
+  int64_t val = ArrowArrayViewGetIntUnsafe(m_array, rowIndex);
+  if (m_useNumpy) {
+    return PyObject_CallMethod(
+        m_context, "INTERVAL_YEAR_MONTH_to_numpy_timedelta", "L", val);
+  }
+  // Python timedelta does not support year-month intervals. Use long instead.
+  return PyLong_FromLongLong(val);
+}
+
+IntervalDayTimeConverterInt::IntervalDayTimeConverterInt(ArrowArrayView* array,
+                                                         PyObject* context,
+                                                         bool useNumpy)
+    : m_array(array), m_context(context) {
+  m_method = useNumpy ? INTERVAL_DT_INT_TO_NUMPY_TIMEDELTA
+                      : INTERVAL_DT_INT_TO_TIMEDELTA;
+}
+
+PyObject* IntervalDayTimeConverterInt::toPyObject(int64_t rowIndex) const {
+  if (ArrowArrayViewIsNull(m_array, rowIndex)) {
+    Py_RETURN_NONE;
+  }
+  int64_t val = ArrowArrayViewGetIntUnsafe(m_array, rowIndex);
+  return PyObject_CallMethod(m_context, m_method, "L", val);
+}
+
+IntervalDayTimeConverterDecimal::IntervalDayTimeConverterDecimal(
+    ArrowArrayView* array, PyObject* context, bool useNumpy)
+    : m_array(array), m_context(context) {
+  m_method = useNumpy ? INTERVAL_DT_DECIMAL_TO_NUMPY_TIMEDELTA
+                      : INTERVAL_DT_DECIMAL_TO_TIMEDELTA;
+}
+
+PyObject* IntervalDayTimeConverterDecimal::toPyObject(int64_t rowIndex) const {
+  if (ArrowArrayViewIsNull(m_array, rowIndex)) {
+    Py_RETURN_NONE;
+  }
+  int64_t bytes_start = 16 * (m_array->array->offset + rowIndex);
+  const char* ptr_start = m_array->buffer_views[1].data.as_char;
+  PyObject* int128_bytes =
+      PyBytes_FromStringAndSize(&(ptr_start[bytes_start]), 16);
+  return PyObject_CallMethod(m_context, m_method, "S", int128_bytes);
+}
+}  // namespace sf

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/IntervalConverter.hpp

@@ -0,0 +1,56 @@
+#ifndef PC_INTERVALCONVERTER_HPP
+#define PC_INTERVALCONVERTER_HPP
+
+#include <memory>
+
+#include "IColumnConverter.hpp"
+#include "nanoarrow.h"
+#include "nanoarrow.hpp"
+
+namespace sf {
+
+class IntervalYearMonthConverter : public IColumnConverter {
+ public:
+  explicit IntervalYearMonthConverter(ArrowArrayView* array, PyObject* context,
+                                      bool useNumpy);
+  virtual ~IntervalYearMonthConverter() = default;
+
+  PyObject* toPyObject(int64_t rowIndex) const override;
+
+ private:
+  ArrowArrayView* m_array;
+  PyObject* m_context;
+  bool m_useNumpy;
+};
+
+class IntervalDayTimeConverterInt : public IColumnConverter {
+ public:
+  explicit IntervalDayTimeConverterInt(ArrowArrayView* array, PyObject* context,
+                                       bool useNumpy);
+  virtual ~IntervalDayTimeConverterInt() = default;
+
+  PyObject* toPyObject(int64_t rowIndex) const override;
+
+ private:
+  ArrowArrayView* m_array;
+  PyObject* m_context;
+  const char* m_method;
+};
+
+class IntervalDayTimeConverterDecimal : public IColumnConverter {
+ public:
+  explicit IntervalDayTimeConverterDecimal(ArrowArrayView* array,
+                                           PyObject* context, bool useNumpy);
+  virtual ~IntervalDayTimeConverterDecimal() = default;
+
+  PyObject* toPyObject(int64_t rowIndex) const override;
+
+ private:
+  ArrowArrayView* m_array;
+  PyObject* m_context;
+  const char* m_method;
+};
+
+}  // namespace sf
+
+#endif  // PC_INTERVALCONVERTER_HPP

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/SnowflakeType.cpp

@@ -15,6 +15,8 @@ std::unordered_map<std::string, SnowflakeType::Type>
         {"FIXED", SnowflakeType::Type::FIXED},
         {"DECFLOAT", SnowflakeType::Type::DECFLOAT},
         {"FLOAT", SnowflakeType::Type::REAL},
+        {"INTERVAL_YEAR_MONTH", SnowflakeType::Type::INTERVAL_YEAR_MONTH},
+        {"INTERVAL_DAY_TIME", SnowflakeType::Type::INTERVAL_DAY_TIME},
         {"MAP", SnowflakeType::Type::MAP},
         {"OBJECT", SnowflakeType::Type::OBJECT},
         {"REAL", SnowflakeType::Type::REAL},

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/SnowflakeType.hpp

@@ -30,6 +30,8 @@ class SnowflakeType {
     VECTOR = 16,
     MAP = 17,
     DECFLOAT = 18,
+    INTERVAL_YEAR_MONTH = 19,
+    INTERVAL_DAY_TIME = 20,
   };
 
   static SnowflakeType::Type snowflakeTypeFromString(std::string str) {

test/integ/test_arrow_result.py

@@ -1223,6 +1223,65 @@ def test_fetch_as_numpy_val(conn_cnx):
         assert val[3] == numpy.datetime64("2019-01-02 12:34:56.12345678")
 
 
+@pytest.mark.parametrize("use_numpy", [True, False])
+def test_select_year_month_interval_arrow(conn_cnx, use_numpy):
+    cases = ["0-0", "1-2", "-1-3", "999999999-11", "-999999999-11"]
+    expected = [0, 14, -15, 11_999_999_999, -11_999_999_999]
+    if use_numpy:
+        expected = [numpy.timedelta64(e, "M") for e in expected]
+
+    table = "test_arrow_day_time_interval"
+    values = "(" + "),(".join([f"'{c}'" for c in cases]) + ")"
+    with conn_cnx(numpy=use_numpy) as conn:
+        cursor = conn.cursor()
+        cursor.execute("alter session set python_connector_query_result_format='arrow'")
+
+        cursor.execute("alter session set feature_interval_types=enabled")
+        cursor.execute(f"create or replace table {table} (c1 interval year to month)")
+        cursor.execute(f"insert into {table} values {values}")
+        result = conn.cursor().execute(f"select * from {table}").fetchall()
+        result = [r[0] for r in result]
+        assert result == expected
+
+
+@pytest.mark.skip(
+    reason="SNOW-1878635: Add support for day-time interval in ArrowStreamWriter"
+)
+@pytest.mark.parametrize("use_numpy", [True, False])
+def test_select_day_time_interval_arrow(conn_cnx, use_numpy):
+    cases = [
+        "0 0:0:0.0",
+        "12 3:4:5.678",
+        "-1 2:3:4.567",
+        "99999 23:59:59.999999",
+        "-99999 23:59:59.999999",
+    ]
+    expected = [
+        timedelta(days=0),
+        timedelta(days=12, hours=3, minutes=4, seconds=5.678),
+        -timedelta(days=1, hours=2, minutes=3, seconds=4.567),
+        timedelta(days=99999, hours=23, minutes=59, seconds=59.999999),
+        -timedelta(days=99999, hours=23, minutes=59, seconds=59.999999),
+    ]
+    if use_numpy:
+        expected = [numpy.timedelta64(e) for e in expected]
+
+    table = "test_arrow_day_time_interval"
+    values = "(" + "),(".join([f"'{c}'" for c in cases]) + ")"
+    with conn_cnx(numpy=use_numpy) as conn:
+        cursor = conn.cursor()
+        cursor.execute("alter session set python_connector_query_result_format='arrow'")
+
+        cursor.execute("alter session set feature_interval_types=enabled")
+        cursor.execute(
+            f"create or replace table {table} (c1 interval day(5) to second)"
+        )
+        cursor.execute(f"insert into {table} values {values}")
+        result = conn.cursor().execute(f"select * from {table}").fetchall()
+        result = [r[0] for r in result]
+        assert result == expected
+
+
 def get_random_seed():
     random.seed(datetime.now().timestamp())
     return random.randint(0, 10000)

test/unit/test_converter.py

@@ -1,9 +1,11 @@
 #!/usr/bin/env python
 from __future__ import annotations
 
+from datetime import timedelta
 from decimal import Decimal
 from logging import getLogger
 
+import numpy
 import pytest
 
 from snowflake.connector import ProgrammingError
@@ -97,3 +99,37 @@ def test_converter_to_snowflake_bindings_error():
         match=r"Binding data in type \(somethingsomething\) is not supported",
     ):
         converter._somethingsomething_to_snowflake_bindings("Bogus")
+
+
+NANOS_PER_DAY = 24 * 60 * 60 * 10**9
+
+
+@pytest.mark.parametrize("nanos", [0, 1, 999, 1000, 999999, 10**5 * NANOS_PER_DAY - 1])
+def test_day_time_interval_int_to_timedelta(nanos):
+    converter = ArrowConverterContext()
+    assert converter.INTERVAL_DAY_TIME_int_to_timedelta(nanos) == timedelta(
+        microseconds=nanos // 1000
+    )
+    assert converter.INTERVAL_DAY_TIME_int_to_numpy_timedelta(
+        nanos
+    ) == numpy.timedelta64(nanos, "ns")
+
+
+@pytest.mark.parametrize("nanos", [0, 1, 999, 1000, 999999, 10**9 * NANOS_PER_DAY - 1])
+def test_day_time_interval_decimal_to_timedelta(nanos):
+    converter = ArrowConverterContext()
+    nano_bytes = nanos.to_bytes(16, byteorder="little", signed=True)
+    assert converter.INTERVAL_DAY_TIME_decimal_to_timedelta(nano_bytes) == timedelta(
+        microseconds=nanos // 1000
+    )
+    assert converter.INTERVAL_DAY_TIME_decimal_to_numpy_timedelta(
+        nano_bytes
+    ) == numpy.timedelta64(nanos // 1_000_000, "ms")
+
+
+@pytest.mark.parametrize("months", [0, 1, 999, 1000, 999999, 10**9 * 12 - 1])
+def test_year_month_interval_to_timedelta(months):
+    converter = ArrowConverterContext()
+    assert converter.INTERVAL_YEAR_MONTH_to_numpy_timedelta(
+        months
+    ) == numpy.timedelta64(months, "M")