add decfloat type

Author: sfc-gh-bkogan

src/snowflake/snowpark/_internal/analyzer/datatype_mapper.py

@@ -25,6 +25,7 @@
     DataType,
     DateType,
     DayTimeIntervalType,
+    DecFloatType,
     DecimalType,
     DoubleType,
     FileType,
@@ -362,6 +363,9 @@ def to_sql(
     if isinstance(datatype, (FloatType, DoubleType)):
         if value is None:
             return "NULL :: FLOAT"
+    if isinstance(datatype, DecFloatType):
+        if value is None:
+            return "NULL :: DECFLOAT"
     if isinstance(datatype, StringType):
         if value is None:
             return f"NULL :: {analyzer_utils.string(datatype.length)}"
@@ -403,6 +407,13 @@ def to_sql(
     if isinstance(datatype, BooleanType):
         return f"{value} :: BOOLEAN"
 
+    # DecFloatType must use DECFLOAT 'value' syntax (not value::DECFLOAT) to preserve precision
+    if isinstance(datatype, DecFloatType):
+        if isinstance(value, str):
+            return f"DECFLOAT {str_to_sql(value)}"
+        else:
+            return f"DECFLOAT '{value}'"
+
     if isinstance(value, float) and isinstance(datatype, _FractionalType):
         if math.isnan(value) or math.isinf(value):
             return float_nan_inf_to_sql(value)
@@ -495,6 +506,8 @@ def schema_expression(data_type: DataType, is_nullable: bool) -> str:
             return "PARSE_JSON('NULL') :: VARIANT"
         return "NULL :: " + convert_sp_to_sf_type(data_type)
 
+    if isinstance(data_type, DecFloatType):
+        return "DECFLOAT '0'"
     if isinstance(data_type, _NumericType):
         return "0 :: " + convert_sp_to_sf_type(data_type)
     if isinstance(data_type, StringType):
@@ -596,6 +609,13 @@ def numeric_to_sql_without_cast(value: Any, datatype: DataType) -> str:
         # regular to_sql generation
         return to_sql(value, datatype)
 
+    # DecFloatType must always use DECFLOAT 'value' syntax to preserve precision
+    if isinstance(datatype, DecFloatType):
+        if isinstance(value, str):
+            return f"DECFLOAT {str_to_sql(value)}"
+        else:
+            return f"DECFLOAT '{value}'"
+
     if isinstance(value, float) and isinstance(datatype, _FractionalType):
         # when the float value is NAN or INF, a cast is still required
         if math.isnan(value) or math.isinf(value):

src/snowflake/snowpark/_internal/type_utils.py

@@ -51,6 +51,7 @@
     DateType,
     DayTimeInterval,
     DayTimeIntervalType,
+    DecFloatType,
     DecimalType,
     DoubleType,
     FloatType,
@@ -188,7 +189,7 @@ def convert_metadata_to_sp_type(
         return convert_sf_to_sp_type(
             column_type_name,
             metadata.precision or 0,
-            metadata.scale or 0,
+            metadata.scale,
             metadata.internal_size or 0,
             max_string_size,
         )
@@ -197,7 +198,7 @@ def convert_metadata_to_sp_type(
 def convert_sf_to_sp_type(
     column_type_name: str,
     precision: int,
-    scale: int,
+    scale: int | None,
     internal_size: int,
     max_string_size: int,
 ) -> DataType:
@@ -291,6 +292,8 @@ def convert_sf_to_sp_type(
         return TimestampType(timezone=TimestampTimeZone.TZ)
     if column_type_name == "DATE":
         return DateType()
+    if column_type_name == "FIXED" and scale is None:
+        return DecFloatType()
     if column_type_name == "DECIMAL" or (
         (column_type_name == "FIXED" or column_type_name == "NUMBER") and scale != 0
     ):
@@ -333,6 +336,8 @@ def convert_sp_to_sf_type(datatype: DataType, nullable_override=None) -> str:
         return "FLOAT"
     if isinstance(datatype, DoubleType):
         return "DOUBLE"
+    if isinstance(datatype, DecFloatType):
+        return "DECFLOAT"
     # We regard NullType as String, which is required when creating
     # a dataframe from local data with all None values
     if isinstance(datatype, StringType):
@@ -845,6 +850,7 @@ def snow_type_to_dtype_str(snow_type: DataType) -> str:
             BooleanType,
             FloatType,
             DoubleType,
+            DecFloatType,
             DateType,
             TimestampType,
             TimeType,

src/snowflake/snowpark/functions.py

@@ -2118,6 +2118,31 @@ def to_double(
     )
 
 
+@publicapi
+def to_decfloat(
+    e: ColumnOrName, fmt: Optional[ColumnOrLiteralStr] = None, _emit_ast: bool = True
+) -> Column:
+    """Converts an input expression to a decimal floating-point number.
+
+    Example::
+        >>> df = session.create_dataframe(['12', '11.3', '-90.12345'], schema=['a'])
+        >>> df.select(to_decfloat(col('a')).as_('ans')).collect()
+        [Row(ANS=Decimal('12.0')), Row(ANS=Decimal('11.3')), Row(ANS=Decimal('-90.12345'))]
+    """
+    ast = (
+        build_function_expr("to_decfloat", [e] if fmt is None else [e, fmt])
+        if _emit_ast
+        else None
+    )
+    c = _to_col_if_str(e, "to_decfloat")
+    fmt_col = _to_col_if_lit(fmt, "to_decfloat") if fmt is not None else None
+    return (
+        _call_function("to_decfloat", c, _ast=ast, _emit_ast=_emit_ast)
+        if fmt_col is None
+        else _call_function("to_decfloat", c, fmt_col, _ast=ast, _emit_ast=_emit_ast)
+    )
+
+
 @publicapi
 def div0(
     dividend: Union[ColumnOrName, int, float],

src/snowflake/snowpark/mock/_snowflake_data_type.py

@@ -11,6 +11,7 @@
     BooleanType,
     DataType,
     DateType,
+    DecFloatType,
     DecimalType,
     DoubleType,
     FloatType,
@@ -114,20 +115,27 @@ class SnowDataTypeConversion:
     SnowDataTypeConversion(BooleanType, DecimalType, True, False),
     SnowDataTypeConversion(BooleanType, StringType, True, True),
     SnowDataTypeConversion(BooleanType, VariantType, True, True),
+    SnowDataTypeConversion(BooleanType, DecFloatType, True, True),
     SnowDataTypeConversion(DateType, TimestampType, True, False),
     SnowDataTypeConversion(DateType, StringType, True, True),
     SnowDataTypeConversion(DateType, VariantType, True, False),
     SnowDataTypeConversion(FloatType, BooleanType, True, True),
     SnowDataTypeConversion(FloatType, DecimalType, True, True),
     SnowDataTypeConversion(FloatType, StringType, True, True),
     SnowDataTypeConversion(FloatType, VariantType, True, True),
+    SnowDataTypeConversion(FloatType, DecFloatType, True, True),
+    SnowDataTypeConversion(DecFloatType, BooleanType, True, True),
+    SnowDataTypeConversion(DecFloatType, DecimalType, True, True),
+    SnowDataTypeConversion(DecFloatType, FloatType, True, True),
+    SnowDataTypeConversion(DecFloatType, StringType, True, True),
     SnowDataTypeConversion(GeographyType, VariantType, True, False),
     # SnowDataTypeConversion(GeometryType, VariantType, True, False),  # GeometryType isn't available yet.
     SnowDataTypeConversion(DecimalType, BooleanType, True, True),
     SnowDataTypeConversion(DecimalType, FloatType, True, True),
     SnowDataTypeConversion(DecimalType, TimestampType, True, True),
     SnowDataTypeConversion(DecimalType, StringType, True, True),
     SnowDataTypeConversion(DecimalType, VariantType, True, True),
+    SnowDataTypeConversion(DecimalType, DecFloatType, True, True),
     SnowDataTypeConversion(MapType, ArrayType, True, False),
     SnowDataTypeConversion(MapType, StringType, True, False),
     SnowDataTypeConversion(MapType, VariantType, True, True),
@@ -140,6 +148,7 @@ class SnowDataTypeConversion:
     SnowDataTypeConversion(StringType, BooleanType, True, True),
     SnowDataTypeConversion(StringType, DateType, True, True),
     SnowDataTypeConversion(StringType, FloatType, True, True),
+    SnowDataTypeConversion(StringType, DecFloatType, True, True),
     SnowDataTypeConversion(StringType, DecimalType, True, True),
     SnowDataTypeConversion(StringType, TimeType, True, True),
     SnowDataTypeConversion(StringType, TimestampType, True, True),
@@ -271,6 +280,9 @@ def calculate_type(c1: ColumnType, c2: ColumnType, op: Union[str]):
                 parameters_info={"op": op},
                 raise_error=NotImplementedError,
             )
+    # DECFLOAT has highest numeric priority
+    elif isinstance(t1, DecFloatType) or isinstance(t2, DecFloatType):
+        return ColumnType(DecFloatType(), nullable)
     elif isinstance(t1, (FloatType, DoubleType)) or isinstance(
         t2, (FloatType, DoubleType)
     ):

src/snowflake/snowpark/mock/_snowflake_to_pandas_converter.py

@@ -21,6 +21,7 @@
     ByteType,
     DataType,
     DateType,
+    DecFloatType,
     DecimalType,
     DoubleType,
     FloatType,
@@ -97,6 +98,19 @@ def _decimal_converter(
         )
 
 
+def _decfloat_converter(
+    value: str, datatype: DataType, field_optionally_enclosed_by=None, null_if=None
+):
+    if value is None or value == "" or (null_if is not None and value in null_if):
+        return None
+    try:
+        return Decimal(value)
+    except Exception as exc:
+        SnowparkLocalTestingException.raise_from_error(
+            exc, error_message=f"Numeric value '{value}' is not recognized."
+        )
+
+
 def _bool_converter(
     value: str,
     datatype: DataType,
@@ -189,6 +203,7 @@ def _time_converter(
     ShortType: _integer_converter,
     DoubleType: _fraction_converter,
     FloatType: _fraction_converter,
+    DecFloatType: _decfloat_converter,
     DecimalType: _decimal_converter,
     BooleanType: _bool_converter,
     DateType: _date_converter,

src/snowflake/snowpark/mock/_util.py

@@ -19,6 +19,7 @@
     BooleanType,
     ByteType,
     DateType,
+    DecFloatType,
     DecimalType,
     DoubleType,
     FloatType,
@@ -244,6 +245,7 @@ def fix_drift_between_column_sf_type_and_dtype(col: ColumnEmulator):
         BooleanType: bool,
         ByteType: numpy.int8 if not col.sf_type.nullable else "Int8",
         DateType: object,
+        DecFloatType: numpy.float64,
         DecimalType: numpy.float64,
         DoubleType: numpy.float64,
         FloatType: numpy.float64,

src/snowflake/snowpark/types.py

@@ -461,6 +461,13 @@ def _fill_ast(self, ast: proto.DataType) -> None:
         ast.double_type = True
 
 
+class DecFloatType(_FractionalType):
+    """DecFloat data type. This maps to the DECFLOAT data type in Snowflake."""
+
+    def _fill_ast(self, ast: proto.DataType) -> None:
+        ast.decfloat_type = True
+
+
 class DecimalType(_FractionalType):
     """Decimal data type. This maps to the NUMBER data type in Snowflake."""
 
@@ -1108,6 +1115,7 @@ def _fill_ast(self, ast: proto.DataType) -> None:
     BinaryType,
     BooleanType,
     DecimalType,
+    DecFloatType,
     FloatType,
     DoubleType,
     ByteType,

tests/integ/test_datatypes.py

@@ -13,6 +13,7 @@
 from snowflake.snowpark.functions import lit
 from snowflake.snowpark.types import (
     BooleanType,
+    DecFloatType,
     DecimalType,
     DoubleType,
     FloatType,
@@ -51,13 +52,14 @@ def test_basic_filter(session):
 
 def test_plus_basic(session):
     df = session.create_dataframe(
-        [[1, 1.1, 2.2, 3.3]],
+        [[1, 1.1, 2.2, 3.3, 4.4]],
         schema=StructType(
             [
                 StructField("a", LongType(), nullable=False),
                 StructField("b", DecimalType(3, 1), nullable=False),
                 StructField("c", DoubleType(), nullable=False),
                 StructField("d", DecimalType(4, 2), nullable=False),
+                StructField("e", DecFloatType(), nullable=False),
             ]
         ),
     )
@@ -66,27 +68,30 @@ def test_plus_basic(session):
         (df["a"] + 1).as_("new_a"),
         (df["b"] + df["d"]).as_("new_b"),
         (df["c"] + 3).as_("new_c"),
+        (df["e"] + df["c"]).as_("new_e"),
     )
     assert repr(df.schema) == repr(
         StructType(
             [
                 StructField("NEW_A", LongType(), nullable=False),
                 StructField("NEW_B", DecimalType(5, 2), nullable=False),
                 StructField("NEW_C", DoubleType(), nullable=False),
+                StructField("NEW_E", DecFloatType(), nullable=False),
             ]
         )
     )
 
 
 def test_minus_basic(session):
     df = session.create_dataframe(
-        [[1, 1.1, 2.2, 3.3]],
+        [[1, 1.1, 2.2, 3.3, 4.4]],
         schema=StructType(
             [
                 StructField("a", LongType(), nullable=False),
                 StructField("b", DecimalType(3, 1), nullable=False),
                 StructField("c", DoubleType(), nullable=False),
                 StructField("d", DecimalType(4, 2), nullable=False),
+                StructField("e", DecFloatType(), nullable=False),
             ]
         ),
     )
@@ -95,27 +100,30 @@ def test_minus_basic(session):
         (df["a"] - 1).as_("new_a"),
         (df["b"] - df["d"]).as_("new_b"),
         (df["c"] - 3).as_("new_c"),
+        (df["e"] - df["a"]).as_("new_e"),
     )
     assert repr(df.schema) == repr(
         StructType(
             [
                 StructField("NEW_A", LongType(), nullable=False),
                 StructField("NEW_B", DecimalType(5, 2), nullable=False),
                 StructField("NEW_C", DoubleType(), nullable=False),
+                StructField("NEW_E", DecFloatType(), nullable=False),
             ]
         )
     )
 
 
 def test_multiple_basic(session):
     df = session.create_dataframe(
-        [[1, 1.1, 2.2, 3.3]],
+        [[1, 1.1, 2.2, 3.3, 4.4]],
         schema=StructType(
             [
                 StructField("a", LongType(), nullable=False),
                 StructField("b", DecimalType(3, 1), nullable=False),
                 StructField("c", FloatType(), nullable=False),
                 StructField("d", DecimalType(4, 2), nullable=False),
+                StructField("e", DecFloatType(), nullable=False),
             ]
         ),
     )
@@ -124,13 +132,15 @@ def test_multiple_basic(session):
         (df["a"] * 1).as_("new_a"),
         (df["b"] * df["d"]).as_("new_b"),
         (df["c"] * 3).as_("new_c"),
+        (df["e"] * df["b"]).as_("new_e"),
     )
     assert repr(df.schema) == repr(
         StructType(
             [
                 StructField("NEW_A", LongType(), nullable=False),
                 StructField("NEW_B", DecimalType(7, 3), nullable=False),
                 StructField("NEW_C", DoubleType(), nullable=False),
+                StructField("NEW_E", DecFloatType(), nullable=False),
             ]
         )
     )