[SNOW-2464063] bugfix: Fix string representation for year-month subtypes

Author: sfc-gh-nkumar

src/snowflake/connector/arrow_context.py

@@ -165,10 +165,14 @@ 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:
+    def INTERVAL_YEAR_MONTH_to_str(self, months: int, scale: int) -> str:
+        return interval_year_month_to_string(months, scale)
+
+    def INTERVAL_YEAR_MONTH_to_numpy_timedelta(
+        self, months: int, scale: int
+    ) -> timedelta64:
+        if scale == 1:  # interval year
+            return numpy.timedelta64(months // 12, "Y")
         return numpy.timedelta64(months, "M")
 
     def INTERVAL_DAY_TIME_int_to_numpy_timedelta(self, nanos: int) -> timedelta64:

src/snowflake/connector/converter.py

@@ -357,7 +357,8 @@ def _BOOLEAN_to_python(
         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))
+        scale = ctx["scale"]
+        return lambda v: interval_year_month_to_string(int(v), scale)
 
     def _INTERVAL_YEAR_MONTH_numpy_to_python(self, ctx: dict[str, Any]) -> Callable:
         return lambda v: numpy.timedelta64(int(v), "M")

src/snowflake/connector/interval_util.py

@@ -1,17 +1,26 @@
 #!/usr/bin/env python
 
 
-def interval_year_month_to_string(interval: int) -> str:
+def interval_year_month_to_string(interval: int, scale: int) -> str:
     """Convert a year-month interval to a string.
 
     Args:
-        interval: The year-month interval.
+        interval: The year-month interval value in months.
+        scale: The scale of the interval which represents subtype as follows:
+            0: INTERVAL YEAR TO MONTH
+            1: INTERVAL YEAR
+            2: INTERVAL MONTH
 
     Returns:
         The string representation of the interval.
     """
     sign = "+" if interval >= 0 else "-"
     interval = abs(interval)
+    if scale == 2:  # INTERVAL MONTH
+        return f"{sign}{interval}"
     years = interval // 12
+    if scale == 1:  # INTERVAL YEAR
+        return f"{sign}{years}"
+    # INTERVAL YEAR TO MONTH
     months = interval % 12
     return f"{sign}{years}-{months:02}"

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

@@ -481,8 +481,21 @@ std::shared_ptr<sf::IColumnConverter> getConverterFromSchema(
     }
 
     case SnowflakeType::Type::INTERVAL_YEAR_MONTH: {
+      struct ArrowStringView scaleString = ArrowCharView(nullptr);
+      int scale = 9;
+      if (metadata != nullptr) {
+        returnCode = ArrowMetadataGetValue(metadata, ArrowCharView("scale"),
+                                           &scaleString);
+        SF_CHECK_ARROW_RC_AND_RETURN(
+            returnCode, nullptr,
+            "[Snowflake Exception] error getting 'scale' from "
+            "Arrow metadata, error code: %d",
+            returnCode);
+        scale =
+            std::stoi(std::string(scaleString.data, scaleString.size_bytes));
+      }
       converter = std::make_shared<sf::IntervalYearMonthConverter>(
-          array, context, useNumpy);
+          array, context, useNumpy, scale);
       break;
     }
 

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

@@ -25,8 +25,8 @@ static constexpr char INTERVAL_YEAR_MONTH_TO_STR[] =
 
 IntervalYearMonthConverter::IntervalYearMonthConverter(ArrowArrayView* array,
                                                        PyObject* context,
-                                                       bool useNumpy)
-    : m_array(array), m_context(context) {
+                                                       bool useNumpy, int scale)
+    : m_array(array), m_context(context), m_scale(scale) {
   m_method = useNumpy ? INTERVAL_YEAR_MONTH_TO_NUMPY_TIMEDELTA
                       : INTERVAL_YEAR_MONTH_TO_STR;
 }
@@ -36,7 +36,7 @@ PyObject* IntervalYearMonthConverter::toPyObject(int64_t rowIndex) const {
     Py_RETURN_NONE;
   }
   int64_t val = ArrowArrayViewGetIntUnsafe(m_array, rowIndex);
-  return PyObject_CallMethod(m_context, m_method, "L", val);
+  return PyObject_CallMethod(m_context, m_method, "Li", val, m_scale);
 }
 
 IntervalDayTimeConverterInt::IntervalDayTimeConverterInt(ArrowArrayView* array,

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

@@ -12,14 +12,15 @@ namespace sf {
 class IntervalYearMonthConverter : public IColumnConverter {
  public:
   explicit IntervalYearMonthConverter(ArrowArrayView* array, PyObject* context,
-                                      bool useNumpy);
+                                      bool useNumpy, int scale);
   virtual ~IntervalYearMonthConverter() = default;
 
   PyObject* toPyObject(int64_t rowIndex) const override;
 
  private:
   ArrowArrayView* m_array;
   PyObject* m_context;
+  int m_scale;
   const char* m_method;
 };
 

test/integ/test_interval_types.py

@@ -43,6 +43,40 @@ def test_select_year_month_interval(conn_cnx, use_numpy, result_format):
         assert result == expected
 
 
+@pytest.mark.parametrize("use_numpy", [True, False])
+@pytest.mark.parametrize("result_format", ["json", "arrow"])
+@pytest.mark.parametrize("datatype", ["YEAR", "MONTH"])
+def test_select_year_month_interval_subtypes(
+    conn_cnx, use_numpy, result_format, datatype
+):
+    cases = ["0", "1", "-1", "999999999", "-999999999"]
+    if use_numpy:
+        expected = [numpy.timedelta64(int(c), datatype[0]) for c in cases]
+    else:
+        expected = ["+0", "+1", "-1", "+999999999", "-999999999"]
+
+    table = "test_year_month_interval_subtypes"
+    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 {datatype})")
+        cursor.execute(f"insert into {table} values {values}")
+        result = cursor.execute(f"select * from {table}").fetchall()
+        # Validate column metadata.
+        type_code = cursor._description[0].type_code
+        assert (
+            constants.FIELD_ID_TO_NAME[type_code] == "INTERVAL_YEAR_MONTH"
+        ), f"invalid column type: {type_code}"
+        # Validate column values.
+        result = [r[0] for r in result]
+        assert result == expected
+
+
 @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):