SNOW-2338989: Ensure Arrow to_pandas maps Interval types (#2536)

sfc-gh-nkrishna · sfc-gh-turbaszek · commit 1d2459f91624 · 2025-10-28T15:58:13.000+01:00
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
@@ -13,6 +13,7 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
   - Added the `oauth_credentials_in_body` parameter supporting an option to send the oauth client credentials in the request body
   - Fix retry behavior for `ECONNRESET` error
   - Added an option to exclude `botocore` and `boto3` dependencies by setting `SNOWFLAKE_NO_BOTO` environment variable during installation
+  - Added support for pandas conversion for Day-time and Year-Month Interval types
 
 - v3.17.4(September 22,2025)
   - Added support for intermediate certificates as roots when they are stored in the trust store
diff --git a/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowTableIterator.cpp b/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowTableIterator.cpp
@@ -68,6 +68,7 @@ void CArrowTableIterator::convertIfNeeded(ArrowSchema* columnSchema,
     case SnowflakeType::Type::DATE:
     case SnowflakeType::Type::REAL:
     case SnowflakeType::Type::TEXT:
+    case SnowflakeType::Type::INTERVAL_YEAR_MONTH:
     case SnowflakeType::Type::VARIANT:
     case SnowflakeType::Type::VECTOR: {
       // Do not need to convert
@@ -174,6 +175,24 @@ void CArrowTableIterator::convertIfNeeded(ArrowSchema* columnSchema,
       break;
     }
 
+    case SnowflakeType::Type::INTERVAL_DAY_TIME: {
+      int scale = 9;
+      if (metadata != nullptr) {
+        struct ArrowStringView scaleString = ArrowCharView(nullptr);
+        returnCode = ArrowMetadataGetValue(metadata, ArrowCharView("scale"),
+                                           &scaleString);
+        SF_CHECK_ARROW_RC(returnCode,
+                          "[Snowflake Exception] error getting 'scale' "
+                          "from Arrow metadata, error code: %d",
+                          returnCode);
+        scale =
+            std::stoi(std::string(scaleString.data, scaleString.size_bytes));
+      }
+      convertIntervalDayTimeColumn_nanoarrow(&columnSchemaView, columnArray,
+                                             scale);
+      break;
+    }
+
     case SnowflakeType::Type::TIME: {
       int scale = 9;
       if (metadata != nullptr) {
@@ -503,6 +522,76 @@ void CArrowTableIterator::
   ArrowArrayMove(newArray, columnArray->array);
 }
 
+void CArrowTableIterator::convertIntervalDayTimeColumn_nanoarrow(
+    ArrowSchemaView* field, ArrowArrayView* columnArray, const int scale) {
+  int returnCode = 0;
+  nanoarrow::UniqueSchema newUniqueField;
+  nanoarrow::UniqueArray newUniqueArray;
+  ArrowSchema* newSchema = newUniqueField.get();
+  ArrowArray* newArray = newUniqueArray.get();
+  ArrowError error;
+
+  // create new schema
+  ArrowSchemaInit(newSchema);
+  newSchema->flags &=
+      (field->schema->flags & ARROW_FLAG_NULLABLE);  // map to nullable()
+
+  returnCode = ArrowSchemaSetTypeDateTime(newSchema, NANOARROW_TYPE_DURATION,
+                                          NANOARROW_TIME_UNIT_NANO, NULL);
+  SF_CHECK_ARROW_RC(returnCode,
+                    "[Snowflake Exception] error setting arrow schema type "
+                    "DateTime, error code: %d",
+                    returnCode);
+
+  returnCode = ArrowSchemaSetName(newSchema, field->schema->name);
+  SF_CHECK_ARROW_RC(
+      returnCode,
+      "[Snowflake Exception] error setting schema name, error code: %d",
+      returnCode);
+
+  returnCode = ArrowArrayInitFromSchema(newArray, newSchema, &error);
+  SF_CHECK_ARROW_RC(returnCode,
+                    "[Snowflake Exception] error initializing ArrowArrayView "
+                    "from schema : %s, error code: %d",
+                    ArrowErrorMessage(&error), returnCode);
+
+  returnCode = ArrowArrayStartAppending(newArray);
+  SF_CHECK_ARROW_RC(
+      returnCode,
+      "[Snowflake Exception] error appending arrow array, error code: %d",
+      returnCode);
+
+  for (int64_t rowIdx = 0; rowIdx < columnArray->array->length; rowIdx++) {
+    if (ArrowArrayViewIsNull(columnArray, rowIdx)) {
+      returnCode = ArrowArrayAppendNull(newArray, 1);
+      SF_CHECK_ARROW_RC(returnCode,
+                        "[Snowflake Exception] error appending null to arrow "
+                        "array, error code: %d",
+                        returnCode);
+    } else {
+      ArrowDecimal arrowDecimal;
+      ArrowDecimalInit(&arrowDecimal, 128, 38, 0);
+      ArrowArrayViewGetDecimalUnsafe(columnArray, rowIdx, &arrowDecimal);
+      auto originalVal = ArrowDecimalGetIntUnsafe(&arrowDecimal);
+      returnCode = ArrowArrayAppendInt(newArray, originalVal);
+      SF_CHECK_ARROW_RC(returnCode,
+                        "[Snowflake Exception] error appending int to arrow "
+                        "array, error code: %d",
+                        returnCode);
+    }
+  }
+
+  returnCode = ArrowArrayFinishBuildingDefault(newArray, &error);
+  SF_CHECK_ARROW_RC(returnCode,
+                    "[Snowflake Exception] error finishing building arrow "
+                    "array: %s, error code: %d",
+                    ArrowErrorMessage(&error), returnCode);
+  field->schema->release(field->schema);
+  ArrowSchemaMove(newSchema, field->schema);
+  columnArray->array->release(columnArray->array);
+  ArrowArrayMove(newArray, columnArray->array);
+}
+
 void CArrowTableIterator::convertTimeColumn_nanoarrow(
     ArrowSchemaView* field, ArrowArrayView* columnArray, const int scale) {
   int returnCode = 0;
diff --git a/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowTableIterator.hpp b/src/snowflake/connector/nanoarrow_cpp/ArrowIterator/CArrowTableIterator.hpp
@@ -94,6 +94,14 @@ class CArrowTableIterator : public CArrowIterator {
                                    ArrowArrayView* columnArray,
                                    const int scale);
 
+  /**
+   * convert Snowflake Interval Day-Time column (Arrow int64/decimal128) to
+   * Arrow Duration column
+   */
+  void convertIntervalDayTimeColumn_nanoarrow(ArrowSchemaView* field,
+                                              ArrowArrayView* columnArray,
+                                              const int scale);
+
   /**
    * convert Snowflake TimestampNTZ/TimestampLTZ column to Arrow Timestamp
    * column
diff --git a/test/integ/aio_it/pandas_it/test_arrow_pandas_async.py b/test/integ/aio_it/pandas_it/test_arrow_pandas_async.py
@@ -662,6 +662,26 @@ async def test_vector(conn_cnx, is_public_test):
             await finish(conn, table)
 
 
+@pytest.mark.skipif(
+    not installed_pandas or no_arrow_iterator_ext,
+    reason="arrow_iterator extension is not built, or pandas is missing.",
+)
+async def test_interval_year_month(conn_cnx):
+    cases = ["1-2", "-1-3", "999999999-11", "-999999999-11"]
+    table = "test_arrow_year_month_interval"
+    values = "(" + "),(".join([f"'{c}'" for c in cases]) + ")"
+    async with conn_cnx() as conn:
+        cursor = conn.cursor()
+        await cursor.execute("alter session set feature_interval_types=enabled")
+        await cursor.execute(
+            f"create or replace table {table} (a interval year to month)"
+        )
+        await cursor.execute(f"insert into {table} values {values}")
+        sql_text = f"select a from {table}"
+        await validate_pandas(conn, sql_text, cases, 1, "one", "interval_year_month")
+        await finish(conn, table)
+
+
 async def validate_pandas(
     cnx_table,
     sql,
@@ -740,6 +760,23 @@ async def validate_pandas(
                         c_case = Decimal(cases[i])
                     elif data_type == "date":
                         c_case = datetime.strptime(cases[i], "%Y-%m-%d").date()
+                    elif data_type == "interval_year_month":
+                        year_month_list = cases[i].split("-")
+                        if len(year_month_list) == 2:
+                            c_case = int(year_month_list[0]) * 12 + int(
+                                year_month_list[1]
+                            )
+                        else:
+                            # negative value
+                            c_case = -(
+                                int(year_month_list[1]) * 12 + int(year_month_list[2])
+                            )
+                    elif data_type == "interval_day_time":
+                        timedelta_split_days = cases[i].split(" ")
+                        pandas_timedelta_str = (
+                            timedelta_split_days[0] + " days " + timedelta_split_days[1]
+                        )
+                        c_case = pandas.to_timedelta(pandas_timedelta_str)
                     elif data_type == "time":
                         time_str_len = 8 if scale == 0 else 9 + scale
                         c_case = cases[i].strip()[:time_str_len]
diff --git a/test/integ/pandas_it/test_arrow_pandas.py b/test/integ/pandas_it/test_arrow_pandas.py
@@ -62,6 +62,42 @@ def test_num_one(conn_cnx):
     fetch_pandas(conn_cnx, sql_exec, row_count, col_count, "one")
 
 
+@pytest.mark.skipif(
+    not installed_pandas or no_arrow_iterator_ext,
+    reason="arrow_iterator extension is not built, or pandas is missing.",
+)
+def test_interval_year_month(conn_cnx):
+    cases = ["1-2", "-1-3", "999999999-11", "-999999999-11"]
+    table = "test_arrow_year_month_interval"
+    values = "(" + "),(".join([f"'{c}'" for c in cases]) + ")"
+    with conn_cnx() as conn:
+        cursor = conn.cursor()
+        cursor.execute("alter session set feature_interval_types=enabled")
+        cursor.execute(f"create or replace table {table} (a interval year to month)")
+        cursor.execute(f"insert into {table} values {values}")
+        sql_text = f"select a from {table}"
+        validate_pandas(conn, sql_text, cases, 1, "one", "interval_year_month")
+        finish(conn, table)
+
+
+@pytest.mark.skipif(
+    not installed_pandas or no_arrow_iterator_ext,
+    reason="arrow_iterator extension is not built, or pandas is missing.",
+)
+def test_interval_day_time(conn_cnx):
+    cases = ["106751 23:47:16.854775807", "0 0:0:0.0", "-5 0:0:0.0"]
+    table = "test_arrow_day_time_interval"
+    values = "(" + "),(".join([f"'{c}'" for c in cases]) + ")"
+    with conn_cnx() as conn:
+        cursor = conn.cursor()
+        cursor.execute("alter session set feature_interval_types=enabled")
+        cursor.execute(f"create or replace table {table} (a interval day to second)")
+        cursor.execute(f"insert into {table} values {values}")
+        sql_text = f"select a from {table}"
+        validate_pandas(conn, sql_text, cases, 1, "one", "interval_day_time")
+        finish(conn, table)
+
+
 @pytest.mark.skipif(
     not installed_pandas or no_arrow_iterator_ext,
     reason="arrow_iterator extension is not built, or pandas is missing.",
@@ -734,6 +770,23 @@ def validate_pandas(
                         c_case = Decimal(cases[i])
                     elif data_type == "date":
                         c_case = datetime.strptime(cases[i], "%Y-%m-%d").date()
+                    elif data_type == "interval_year_month":
+                        year_month_list = cases[i].split("-")
+                        if len(year_month_list) == 2:
+                            c_case = int(year_month_list[0]) * 12 + int(
+                                year_month_list[1]
+                            )
+                        else:
+                            # negative value
+                            c_case = -(
+                                int(year_month_list[1]) * 12 + int(year_month_list[2])
+                            )
+                    elif data_type == "interval_day_time":
+                        timedelta_split_days = cases[i].split(" ")
+                        pandas_timedelta_str = (
+                            timedelta_split_days[0] + " days " + timedelta_split_days[1]
+                        )
+                        c_case = pandas.to_timedelta(pandas_timedelta_str)
                     elif data_type == "time":
                         time_str_len = 8 if scale == 0 else 9 + scale
                         c_case = cases[i].strip()[:time_str_len]
