SNOW-2119489: Add support for interval types in json

Author: sfc-gh-nkumar

DESCRIPTION.md

@@ -7,6 +7,9 @@ https://docs.snowflake.com/
 Source code is also available at: https://github.com/snowflakedb/snowflake-connector-python
 
 # Release Notes
+- v3.16.0(TBD)
+  - Added basic json support for Interval types.
+
 - v3.15.1(May 20, 2025)
   - Added basic arrow support for Interval types.
   - Fix `write_pandas` special characters usage in the location name.

src/snowflake/connector/arrow_context.py

@@ -13,6 +13,7 @@
 
 from .constants import PARAMETER_TIMEZONE
 from .converter import _generate_tzinfo_from_tzoffset
+from .interval_util import interval_year_month_to_string
 
 if TYPE_CHECKING:
     from numpy import datetime64, float64, int64, timedelta64
@@ -164,6 +165,9 @@ 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_str(self, months: int) -> str:
+        return interval_year_month_to_string(months)
+
     def INTERVAL_YEAR_MONTH_to_numpy_timedelta(self, months: int) -> timedelta64:
         return numpy.timedelta64(months, "M")
 

src/snowflake/connector/converter.py

@@ -20,6 +20,7 @@
 from .compat import IS_BINARY, IS_NUMERIC
 from .errorcode import ER_NOT_SUPPORT_DATA_TYPE
 from .errors import ProgrammingError
+from .interval_util import interval_year_month_to_string
 from .sfbinaryformat import binary_to_python, binary_to_snowflake
 from .sfdatetime import sfdatetime_total_seconds_from_timedelta
 
@@ -355,6 +356,28 @@ def _BOOLEAN_to_python(
     ) -> Callable:
         return lambda value: value in ("1", "TRUE")
 
+    def _INTERVAL_YEAR_MONTH_to_python(self, ctx: dict[str, Any]) -> Callable:
+        return lambda v: interval_year_month_to_string(int(v))
+
+    def _INTERVAL_YEAR_MONTH_numpy_to_python(self, ctx: dict[str, Any]) -> Callable:
+        return lambda v: numpy.timedelta64(int(v), "M")
+
+    def _INTERVAL_DAY_TIME_to_python(self, ctx: dict[str, Any]) -> Callable:
+        # Python timedelta only supports microsecond precision. We receive value in
+        # nanoseconds.
+        return lambda v: timedelta(microseconds=int(v) // 1000)
+
+    def _INTERVAL_DAY_TIME_numpy_to_python(self, ctx: dict[str, Any]) -> Callable:
+        # Last 4 bits of the precision are used to store the leading field precision of
+        # the interval.
+        lfp = ctx["precision"] & 0x0F
+        # Numpy timedelta only supports up to 64-bit integers. If the leading field
+        # precision is higher than 5 we receive 16 byte integer from server. So we need
+        # to change the unit to milliseconds to fit in 64-bit integer.
+        if lfp > 5:
+            return lambda v: numpy.timedelta64(int(v) // 1_000_000, "ms")
+        return lambda v: numpy.timedelta64(int(v), "ns")
+
     def snowflake_type(self, value: Any) -> str | None:
         """Returns Snowflake data type for the value. This is used for qmark parameter style."""
         type_name = value.__class__.__name__.lower()

src/snowflake/connector/interval_util.py

@@ -0,0 +1,17 @@
+#!/usr/bin/env python
+
+
+def interval_year_month_to_string(interval: int) -> str:
+    """Convert a year-month interval to a string.
+
+    Args:
+        interval: The year-month interval.
+
+    Returns:
+        The string representation of the interval.
+    """
+    sign = "+" if interval >= 0 else "-"
+    interval = abs(interval)
+    years = interval // 12
+    months = interval % 12
+    return f"{sign}{years}-{months:02}"

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

@@ -16,23 +16,27 @@ 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";
+static constexpr char INTERVAL_YEAR_MONTH_TO_NUMPY_TIMEDELTA[] =
+    "INTERVAL_YEAR_MONTH_to_numpy_timedelta";
+// Python timedelta does not support year-month intervals. Use ANSI SQL
+// formatted string instead.
+static constexpr char INTERVAL_YEAR_MONTH_TO_STR[] =
+    "INTERVAL_YEAR_MONTH_to_str";
 
 IntervalYearMonthConverter::IntervalYearMonthConverter(ArrowArrayView* array,
                                                        PyObject* context,
                                                        bool useNumpy)
-    : m_array(array), m_context(context), m_useNumpy(useNumpy) {}
+    : m_array(array), m_context(context) {
+  m_method = useNumpy ? INTERVAL_YEAR_MONTH_TO_NUMPY_TIMEDELTA
+                      : INTERVAL_YEAR_MONTH_TO_STR;
+}
 
 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);
+  return PyObject_CallMethod(m_context, m_method, "L", val);
 }
 
 IntervalDayTimeConverterInt::IntervalDayTimeConverterInt(ArrowArrayView* array,

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

@@ -20,7 +20,7 @@ class IntervalYearMonthConverter : public IColumnConverter {
  private:
   ArrowArrayView* m_array;
   PyObject* m_context;
-  bool m_useNumpy;
+  const char* m_method;
 };
 
 class IntervalDayTimeConverterInt : public IColumnConverter {

test/integ/test_arrow_result.py

@@ -1224,65 +1224,6 @@ 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/integ/test_interval_types.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+from __future__ import annotations
+
+from datetime import timedelta
+
+import numpy
+import pytest
+
+pytestmark = pytest.mark.skipolddriver  # old test driver tests won't run this module
+
+
+@pytest.mark.parametrize("use_numpy", [True, False])
+@pytest.mark.parametrize("result_format", ["json", "arrow"])
+def test_select_year_month_interval(conn_cnx, use_numpy, result_format):
+    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]
+    else:
+        expected = ["+0-00", "+1-02", "-1-03", "+999999999-11", "-999999999-11"]
+
+    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(
+            f"alter session set python_connector_query_result_format='{result_format}'"
+        )
+
+        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])
+@pytest.mark.parametrize("result_format", ["json", "arrow"])
+def test_select_day_time_interval(conn_cnx, use_numpy, result_format):
+    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(
+            f"alter session set python_connector_query_result_format='{result_format}'"
+        )
+
+        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