add default precision

Author: sfc-gh-yuwang

src/snowflake/snowpark/context.py

@@ -43,6 +43,11 @@
 # Global flag for fix 2360274. When enabled schema queries will use NULL as a place holder for any values inside structured objects
 _enable_fix_2360274 = False
 
+# internal only dictionary store the default precision of integral types, if the type does not appear in the
+# dictionary, the default precision is None.
+# example: _integral_type_default_precision = {"IntegerType": 9}, IntegerType default _precision is 9 now
+_integral_type_default_precision = {}
+
 
 def configure_development_features(
     *,

src/snowflake/snowpark/types.py

@@ -371,7 +371,10 @@ def _fill_ast(self, ast: proto.DataType) -> None:
 # Numeric types
 class _IntegralType(_NumericType):
     def __init__(self, **kwargs) -> None:
-        self._precision = kwargs.pop("_precision", None)
+        class_name = type(self).__name__
+        self._precision = kwargs.pop(
+            "_precision", context._integral_type_default_precision.get(class_name, None)
+        )
 
         if kwargs != {}:
             raise TypeError(

tests/integ/test_datatypes.py

@@ -20,6 +20,8 @@
     StringType,
     StructField,
     StructType,
+    IntegerType,
+    ShortType,
 )
 from tests.utils import Utils
 
@@ -524,3 +526,22 @@ def test_write_to_sf_with_correct_precision(session, precision):
         result = session.sql(f"select * from {table_name}")
         datatype = result.schema.fields[0].datatype
         assert datatype._precision == precision
+
+
+@pytest.mark.parametrize(
+    "mock_dict",
+    [
+        {"IntegerType": 5, "LongType": 4},
+        {"LongType": 19, "IntegerType": 10},
+    ],
+)
+def test_integral_type_default_precision(mock_dict):
+    with mock.patch.object(context, "_integral_type_default_precision", mock_dict):
+        integer_type = IntegerType()
+        assert integer_type._precision == mock_dict["IntegerType"]
+
+        long_type = LongType()
+        assert long_type._precision == mock_dict["LongType"]
+
+        short_type = ShortType()
+        assert short_type._precision is None