snowpark-connect 0.28.1__py3-none-any.whl → 0.30.0__py3-none-any.whl

snowflake/snowpark_connect/expression/map_sql_expression.py

@@ -15,6 +15,7 @@ from pyspark.sql.connect import functions as pyspark_functions
 
 import snowflake.snowpark_connect.proto.snowflake_expression_ext_pb2 as snowflake_proto
 from snowflake.snowpark._internal.analyzer.analyzer_utils import unquote_if_quoted
+from snowflake.snowpark.types import DayTimeIntervalType, YearMonthIntervalType
 from snowflake.snowpark_connect.column_name_handler import ColumnNameMap
 from snowflake.snowpark_connect.config import global_config
 from snowflake.snowpark_connect.typed_column import TypedColumn
@@ -33,6 +34,24 @@ from .typer import ExpressionTyper
 
 DECIMAL_RE = re.compile(r"decimal\((\d+), *(\d+)\)")
 
+_INTERVAL_YEARMONTH_PATTERN_RE = re.compile(r"interval (year|month)( to (year|month))?")
+_INTERVAL_DAYTIME_PATTERN_RE = re.compile(
+    r"interval (day|hour|minute|second)( to (day|hour|minute|second))?"
+)
+
+# Interval field mappings using proper constants
+_YEAR_MONTH_FIELD_MAP = {
+    "year": YearMonthIntervalType.YEAR,
+    "month": YearMonthIntervalType.MONTH,
+}
+
+_DAY_TIME_FIELD_MAP = {
+    "day": DayTimeIntervalType.DAY,
+    "hour": DayTimeIntervalType.HOUR,
+    "minute": DayTimeIntervalType.MINUTE,
+    "second": DayTimeIntervalType.SECOND,
+}
+
 _window_specs = ContextVar[dict[str, any]]("_window_specs", default={})
 
 
@@ -388,8 +407,94 @@ def map_logical_plan_expression(exp: jpype.JObject) -> expressions_proto.Express
                 type_value = types_proto.DataType()
             elif type_name == "binary":
                 type_value = bytes(type_value)
-            elif type_name.startswith("interval "):
-                type_name = "day_time_interval"  # TODO
+            elif year_month_match := _INTERVAL_YEARMONTH_PATTERN_RE.match(type_name):
+                # Extract start and end fields for year-month intervals
+                start_field_name = year_month_match.group(1)  # 'year' or 'month'
+                end_field_name = (
+                    year_month_match.group(3)
+                    if year_month_match.group(3)
+                    else start_field_name
+                )
+
+                # Validate field names exist in mapping
+                start_field = _YEAR_MONTH_FIELD_MAP.get(start_field_name)
+                end_field = _YEAR_MONTH_FIELD_MAP.get(end_field_name)
+
+                if start_field is None:
+                    raise AnalysisException(
+                        f"Invalid year-month interval start field: '{start_field_name}'. Expected 'year' or 'month'."
+                    )
+                if end_field is None:
+                    raise AnalysisException(
+                        f"Invalid year-month interval end field: '{end_field_name}'. Expected 'year' or 'month'."
+                    )
+
+                # Validate field ordering (start_field should be <= end_field)
+                if start_field > end_field:
+                    raise AnalysisException(
+                        f"Invalid year-month interval: start field '{start_field_name}' must come before or equal to end field '{end_field_name}'."
+                    )
+
+                # Use extension for year-month intervals to preserve start/end field info
+                literal = expressions_proto.Expression.Literal(
+                    year_month_interval=type_value
+                )
+                any_proto = Any()
+                any_proto.Pack(
+                    snowflake_proto.ExpExtension(
+                        interval_literal=snowflake_proto.IntervalLiteralExpression(
+                            literal=literal,
+                            start_field=start_field,
+                            end_field=end_field,
+                        )
+                    )
+                )
+                return expressions_proto.Expression(extension=any_proto)
+            elif day_time_match := _INTERVAL_DAYTIME_PATTERN_RE.match(type_name):
+                # Extract start and end fields for day-time intervals
+                start_field_name = day_time_match.group(
+                    1
+                )  # 'day', 'hour', 'minute', 'second'
+                end_field_name = (
+                    day_time_match.group(3)
+                    if day_time_match.group(3)
+                    else start_field_name
+                )
+
+                # Validate field names exist in mapping
+                start_field = _DAY_TIME_FIELD_MAP.get(start_field_name)
+                end_field = _DAY_TIME_FIELD_MAP.get(end_field_name)
+
+                if start_field is None:
+                    raise AnalysisException(
+                        f"Invalid day-time interval start field: '{start_field_name}'. Expected 'day', 'hour', 'minute', or 'second'."
+                    )
+                if end_field is None:
+                    raise AnalysisException(
+                        f"Invalid day-time interval end field: '{end_field_name}'. Expected 'day', 'hour', 'minute', or 'second'."
+                    )
+
+                # Validate field ordering (start_field should be <= end_field)
+                if start_field > end_field:
+                    raise AnalysisException(
+                        f"Invalid day-time interval: start field '{start_field_name}' must come before or equal to end field '{end_field_name}'."
+                    )
+
+                # Use extension for day-time intervals to preserve start/end field info
+                literal = expressions_proto.Expression.Literal(
+                    day_time_interval=type_value
+                )
+                any_proto = Any()
+                any_proto.Pack(
+                    snowflake_proto.ExpExtension(
+                        interval_literal=snowflake_proto.IntervalLiteralExpression(
+                            literal=literal,
+                            start_field=start_field,
+                            end_field=end_field,
+                        )
+                    )
+                )
+                return expressions_proto.Expression(extension=any_proto)
             elif m := DECIMAL_RE.fullmatch(type_name):
                 type_name = "decimal"
                 type_value = expressions_proto.Expression.Literal.Decimal(