diff --git a/.github/workflows/precommit.yml b/.github/workflows/precommit.yml index ad5cc50be0..e587461fb9 100644 --- a/.github/workflows/precommit.yml +++ b/.github/workflows/precommit.yml @@ -238,15 +238,21 @@ jobs: .tox/coverage.xml test-local-testing: - name: Test Local Testing Module py-${{ matrix.os }}-${{ matrix.python-version }} + name: Test Local Testing Module py-${{ matrix.os }}-${{ matrix.python-version }}-pandas${{ matrix.pandas-major }} needs: build runs-on: ${{ matrix.os }} strategy: fail-fast: false matrix: - os: [ubuntu-latest] - python-version: ["3.14"] - cloud-provider: [azure] + include: + - os: ubuntu-latest + python-version: "3.14" + cloud-provider: azure + pandas-major: "2" + - os: ubuntu-latest + python-version: "3.14" + cloud-provider: azure + pandas-major: "3" steps: - name: Checkout Code uses: actions/checkout@v4 @@ -286,7 +292,7 @@ jobs: - name: Install tox run: uv pip install tox --system - name: Run tests - run: python -m tox -e "py${PYTHON_VERSION/\./}-local" + run: python -m tox -e "py${PYTHON_VERSION/\./}-local-pandas${{ matrix.pandas-major }}" env: PYTHON_VERSION: ${{ matrix.python-version }} cloud_provider: ${{ matrix.cloud-provider }} @@ -302,7 +308,7 @@ jobs: - uses: actions/upload-artifact@v4 with: include-hidden-files: true - name: coverage_${{ matrix.os }}-${{ matrix.python-version }}-local-testing + name: coverage_${{ matrix.os }}-${{ matrix.python-version }}-pandas${{ matrix.pandas-major }}-local-testing path: | .tox/.coverage .tox/coverage.xml diff --git a/CHANGELOG.md b/CHANGELOG.md index ea750f0639..71647f9bea 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,26 @@ - Added interval type support for Python UDFs and stored procedures. Use `datetime.timedelta` as the type annotation for day-time interval (`DayTimeIntervalType`) parameters and return values, and `YearMonthInterval` (a type annotation sentinel from `snowflake.snowpark.types`) for year-month interval (`YearMonthIntervalType`) parameters and return values. +#### Dependency Updates + +- Added support for pandas 3, lifting the `pandas<3.0.0` cap added to the `[pandas]` extra in 1.53.0. `snowflake-snowpark-python[pandas]` now declares its own `pandas<4.0.0` and `pyarrow` requirements instead of depending on `snowflake-connector-python[pandas]`. pandas 2 remains fully supported. Upgrading Snowpark does not upgrade an existing pandas 2 install: + - pandas 3 requires Python 3.11 or later, so Python 3.10 environments continue to resolve pandas 2. + - `snowflake-snowpark-python[modin]` continues to resolve pandas 2, because the `[modin]` extra declares `pandas<=2.4`. + +#### Behavior Changes + +- On pandas 3, `DataFrame.to_pandas()` and `DataFrame.to_pandas_batches()` return text columns as the `str` dtype instead of `object`, so a SQL NULL in them arrives as `nan` instead of `None`. Test for NULL with `pandas.isna(value)`, not `value is None`. This also covers `VARIANT`, `OBJECT`, `ARRAY`, `MAP` and the geospatial types, and the `future.infer_string` option does not opt out of it. + +#### Bug Fixes + +- Fixed a bug where `Session.write_pandas`, `Session.write_arrow`, and `Session.create_dataframe` stored a `timedelta` at the column's own resolution instead of nanoseconds, so a microsecond column was written 1000 times too small. Duration columns are now converted to nanoseconds first. pandas 2 was affected too, and rows written by an earlier release need correcting. +- Fixed a local testing bug where a SQL NULL became `nan` instead of `None` after a column was rebuilt through pandas, so `IS NULL` evaluated to `False` and `strict=True` UDF handlers received `nan`. This affected `parse_json`, `to_char`, `concat`, `concat_ws`, `strip_null_value`, `initcap`, `greatest`, `least`, and VARIANT/OBJECT/ARRAY/MAP subfield access, on pandas 2 as well as pandas 3. +- The following local testing bugs appeared only when running with pandas 3: + - Fixed a bug where `ORDER BY` treated SQL NULL as `nan`. + - Fixed a bug where columns omitted from a MERGE insert kept `nan` instead of NULL under Copy-on-Write. + - Fixed a bug where `group_by` and `pivot` results lost their column and index names. + - Fixed a bug where `create_dataframe` raised `TypeError` for `VARIANT` columns. + ## 1.54.0 (2026-07-29) ### Snowpark Python API updates diff --git a/setup.py b/setup.py index 25d8590197..298cc162f7 100644 --- a/setup.py +++ b/setup.py @@ -43,8 +43,12 @@ REQUIRED_PYTHON_VERSION = ">=3.10" PANDAS_REQUIREMENTS = [ - f"snowflake-connector-python[pandas]{CONNECTOR_DEPENDENCY_VERSION}", - "pandas<3.0.0", + "pandas<4.0.0", + # Mirrors the pyarrow bounds in snowflake-connector-python's [pandas] extra + # (its setup.cfg). We no longer request that extra, but the connector's Arrow + # code still needs this range and warns at import when it is not met. + "pyarrow>=14.0.1,<24; python_version >= '3.14'", + "pyarrow>=14.0.1; python_version < '3.14'", ] MODIN_REQUIREMENTS = [ *PANDAS_REQUIREMENTS, diff --git a/src/snowflake/snowpark/mock/_connection.py b/src/snowflake/snowpark/mock/_connection.py index 92fea43002..0451462e25 100644 --- a/src/snowflake/snowpark/mock/_connection.py +++ b/src/snowflake/snowpark/mock/_connection.py @@ -38,7 +38,11 @@ from snowflake.snowpark.exceptions import SnowparkSessionException from snowflake.snowpark.mock._options import pandas from snowflake.snowpark.mock._plan import MockExecutionPlan, execute_mock_plan -from snowflake.snowpark.mock._snowflake_data_type import ColumnEmulator, TableEmulator +from snowflake.snowpark.mock._snowflake_data_type import ( + ColumnEmulator, + TableEmulator, + isna_helper, +) from snowflake.snowpark.mock._stage_registry import StageEntityRegistry from snowflake.snowpark.mock._telemetry import LocalTestOOBTelemetryService from snowflake.snowpark.mock._util import get_fully_qualified_name @@ -653,7 +657,7 @@ def execute( from snowflake.snowpark.mock import CUSTOM_JSON_ENCODER for idx, row in res.iterrows(): - if row[col] is not None: + if not isna_helper(row[col]): # Snowflake sorts maps by key before serializing if isinstance(row[col], dict): row[col] = dict(sorted(row[col].items())) diff --git a/src/snowflake/snowpark/mock/_functions.py b/src/snowflake/snowpark/mock/_functions.py index 12e8d23d76..0dfc6a69c5 100644 --- a/src/snowflake/snowpark/mock/_functions.py +++ b/src/snowflake/snowpark/mock/_functions.py @@ -1303,7 +1303,11 @@ def convert_char(data, _fmt): raise_error=NotImplementedError, ) - res = column.combine(fmt, convert_char) + res = ColumnEmulator( + [convert_char(v, f) for v, f in zip(column, fmt)], + index=column.index, + dtype=object, + ).__finalize__(column) res.sf_type = ColumnType(StringType(), column.sf_type.nullable) return res @@ -1553,11 +1557,14 @@ def mock_parse_json(expr: ColumnEmulator): from snowflake.snowpark.mock import CUSTOM_JSON_DECODER if isinstance(expr.sf_type.datatype, StringType): - res = expr.apply( - lambda x: try_convert( - partial(json.loads, cls=CUSTOM_JSON_DECODER), False, x - ) - ) + res = ColumnEmulator( + [ + try_convert(partial(json.loads, cls=CUSTOM_JSON_DECODER), False, x) + for x in expr + ], + index=expr.index, + dtype=object, + ).__finalize__(expr) else: res = expr.copy() res.sf_type = ColumnType(VariantType(), expr.sf_type.nullable) @@ -1597,6 +1604,8 @@ def convert_variant_to_array(val): def mock_strip_null_value(expr: ColumnEmulator): return ColumnEmulator( [None if x == "null" else x for x in expr], + index=expr.index, + dtype=object, sf_type=ColumnType(expr.sf_type.datatype, True), ) @@ -1981,14 +1990,22 @@ def _greatest(x: CompareType, y: Any) -> Union[CompareType, float]: @patch("greatest") def mock_greatest(*exprs: ColumnEmulator): - result = reduce(lambda x, y: x.combine(y, _greatest), exprs) + result = ColumnEmulator( + [reduce(_greatest, row) for row in zip(*exprs)], + index=exprs[0].index, + dtype=object, + ).__finalize__(exprs[0]) result.sf_type = exprs[0].sf_type return result @patch("least") def mock_least(*exprs: ColumnEmulator): - result = reduce(lambda x, y: x.combine(y, _least), exprs) + result = ColumnEmulator( + [reduce(_least, row) for row in zip(*exprs)], + index=exprs[0].index, + dtype=object, + ).__finalize__(exprs[0]) result.sf_type = exprs[0].sf_type return result @@ -2032,8 +2049,16 @@ def _initcap(value: Optional[str], delimiters: Optional[str]) -> str: @patch("initcap") -def mock_initcap(values: ColumnEmulator, delimiters: ColumnEmulator): - result = values.combine(delimiters, _initcap) +def mock_initcap(values: ColumnEmulator, delimiters: ColumnEmulator = None): + if delimiters is None: + data = [_initcap(v, None) for v in values] + else: + data = [_initcap(v, d) for v, d in zip(values, delimiters)] + result = ColumnEmulator( + data, + index=values.index, + dtype=object, + ).__finalize__(values) result.sf_type = values.sf_type return result @@ -2149,8 +2174,14 @@ def mock_concat(*columns: ColumnEmulator) -> ColumnEmulator: ValueError("concat expects one or more column(s) to be passed in.") ) pdf = pandas.concat(columns, axis=1) - result = pdf.T.apply( - lambda c: None if c.isnull().values.any() else c.astype(str).str.cat() + transposed = pdf.T + result = ColumnEmulator( + [ + None if c.isnull().values.any() else c.astype(str).str.cat() + for _, c in transposed.items() + ], + index=transposed.columns, + dtype=object, ) result.sf_type = ColumnType(StringType(), result.hasnans) return result @@ -2165,10 +2196,14 @@ def mock_concat_ws(*columns: ColumnEmulator) -> ColumnEmulator: ) ) pdf = pandas.concat(columns, axis=1) - result = pdf.T.apply( - lambda c: None - if c.isnull().values.any() - else c[1:].astype(str).str.cat(sep=c[0]) + transposed = pdf.T + result = ColumnEmulator( + [ + None if c.isnull().values.any() else c[1:].astype(str).str.cat(sep=c[0]) + for _, c in transposed.items() + ], + index=transposed.columns, + dtype=object, ) result.sf_type = ColumnType(StringType(), result.hasnans) return result diff --git a/src/snowflake/snowpark/mock/_plan.py b/src/snowflake/snowpark/mock/_plan.py index f794f4f72f..0c2b549545 100644 --- a/src/snowflake/snowpark/mock/_plan.py +++ b/src/snowflake/snowpark/mock/_plan.py @@ -2005,9 +2005,9 @@ def flatten_object_cell_func(cell): for unspecified_col in set(rows_to_insert.columns).difference( inserted_columns ): - rows_to_insert[unspecified_col].replace( - np.nan, None, inplace=True - ) + rows_to_insert[unspecified_col] = rows_to_insert[ + unspecified_col + ].replace(np.nan, None) else: if len(clause.values) != len(rows_to_insert.columns): @@ -2966,14 +2966,20 @@ def get_bound(bound): and field in col[index] and col[index][field] is None ] - res = col.apply(lambda x: None if x is None or field not in x else x[field]) + res = ColumnEmulator( + [None if x is None or field not in x else x[field] for x in col], + index=col.index, + dtype=object, + ).__finalize__(col) res.sf_type = ColumnType(VariantType(), col.sf_type.nullable) return res elif isinstance(exp, SubfieldInt): col = calculate_expression(exp.child, input_data, analyzer, expr_to_alias) - res = col.apply( - lambda x: None if x is None or exp.field >= len(x) else x[exp.field] - ) + res = ColumnEmulator( + [None if x is None or exp.field >= len(x) else x[exp.field] for x in col], + index=col.index, + dtype=object, + ).__finalize__(col) res.sf_type = ColumnType(VariantType(), col.sf_type.nullable) return res elif isinstance(exp, SnowflakeUDF): diff --git a/src/snowflake/snowpark/mock/_snowflake_data_type.py b/src/snowflake/snowpark/mock/_snowflake_data_type.py index b588de7b2e..5686d97561 100644 --- a/src/snowflake/snowpark/mock/_snowflake_data_type.py +++ b/src/snowflake/snowpark/mock/_snowflake_data_type.py @@ -446,7 +446,39 @@ def __setitem__(self, key, value): self.sf_types[key] = value.sf_type self._null_rows_idxs_map[key] = value._null_rows_idxs + def _object_dtype_sort_key(self, key, by): + """Wrap a ``sort_values(key=...)`` callable so it sees the untouched column. + + pandas re-wraps each sort column as ``pandas.Series(ndarray)``, which under + pandas 3 infers ``str`` dtype, whose NA sentinel is ``nan`` -- so a NULL + stored as ``None`` reaches the comparator as ``nan``. Restore ``object`` + dtype here rather than in the comparator: ``NaN`` sorts largest while NULL + obeys ``nulls_first``. + """ + labels = list(by) if isinstance(by, (list, tuple)) else [by] + + def wrapper(series): + # One sort column means we know the label outright; only a multi-column + # sort has to fall back on the name pandas set on the re-wrapped Series. + label = labels[0] if len(labels) == 1 else getattr(series, "name", None) + if label is not None: + try: + original = PandasDataframeType.__getitem__(self, label) + except (KeyError, IndexError): + original = None + if isinstance(original, PandasSeriesType) and len(original) == len( + series + ): + series = pd.Series( + original.to_numpy(dtype=object), name=label, dtype=object + ) + return key(series) + + return wrapper + def sort_values(self, by, **kwargs): + if kwargs.get("key") is not None and kwargs.get("axis", 0) in (0, "index"): + kwargs["key"] = self._object_dtype_sort_key(kwargs["key"], by) result = super().sort_values(by, **kwargs) result.sf_types = self.sf_types return result @@ -506,7 +538,7 @@ def broadcast_value(value: Any, len: int) -> "ColumnEmulator": class ColumnEmulator(PandasSeriesType): - _metadata = ["sf_type", "_null_rows_idxs"] + _metadata = ["_name", "sf_type", "_null_rows_idxs"] @property def _constructor(self): diff --git a/src/snowflake/snowpark/session.py b/src/snowflake/snowpark/session.py index da0f4a409f..27228b25c9 100644 --- a/src/snowflake/snowpark/session.py +++ b/src/snowflake/snowpark/session.py @@ -3234,6 +3234,29 @@ def get_session_stage( self._session_stage = full_qualified_stage_name return f"{STAGE_PREFIX}{self._session_stage}" + @staticmethod + def _normalize_arrow_duration_to_ns(table: "pyarrow.Table") -> "pyarrow.Table": + # Snowflake stores duration as an integer with no unit. Cast to ns + # so a us column is not written 1000x too small. + arrays = [] + fields = [] + changed = False + for i in range(table.num_columns): + field = table.schema.field(i) + column = table.column(i) + if pyarrow.types.is_duration(field.type) and field.type.unit != "ns": + column = column.cast(pyarrow.duration("ns")) + field = field.with_type(pyarrow.duration("ns")) + changed = True + arrays.append(column) + fields.append(field) + if not changed: + return table + # Carry the original schema metadata over: pyarrow.schema() drops it otherwise. + return pyarrow.Table.from_arrays( + arrays, schema=pyarrow.schema(fields, metadata=table.schema.metadata) + ) + @experimental(version="1.28.0") @publicapi def write_arrow( @@ -3295,6 +3318,7 @@ def write_arrow( set use_logical_type as True. Set to None to use Snowflakes default. For more information, see: https://docs.snowflake.com/en/sql-reference/sql/create-file-format """ + table = self._normalize_arrow_duration_to_ns(table) cursor = self._conn._conn.cursor() if quote_identifiers: @@ -3610,6 +3634,19 @@ def write_pandas( # TODO: Implement here write_pandas correctly. success, ci_output = True, [] else: + # write_pandas stores timedelta as a NUMBER of ticks with + # no unit. Normalize to ns so pandas 3's us default does + # not shrink values 1000x versus existing tables. + timedelta_columns = [ + td_col + for td_col in df.columns + if pandas.api.types.is_timedelta64_dtype(df[td_col].dtype) + and str(df[td_col].dtype) != "timedelta64[ns]" + ] + if timedelta_columns: + df = df.copy(deep=False) + for td_col in timedelta_columns: + df[td_col] = df[td_col].astype("timedelta64[ns]") success, _, _, ci_output = write_pandas( self._conn._conn, df, diff --git a/tests/integ/scala/test_datatype_suite.py b/tests/integ/scala/test_datatype_suite.py index 950e05a060..05c38791ab 100644 --- a/tests/integ/scala/test_datatype_suite.py +++ b/tests/integ/scala/test_datatype_suite.py @@ -504,13 +504,22 @@ def test_structured_dtypes_select( reason="FEAT: SNOW-1372813 Cast to StructType not supported", ) def test_structured_dtypes_pandas(structured_type_session, structured_type_support): + # Not imported from _internal.utils: stored procs run the server's released snowpark. + import pandas + + pandas_major_version = int(pandas.__version__.split(".")[0]) + pdf = _create_test_dataframe( structured_type_session, structured_type_support ).to_pandas() if structured_type_support: - assert ( - pdf.to_json() - == '{"MAP":{"0":[["k1",1.0]]},"OBJ":{"0":{"A":"foo","b":0.05}},"ARR":{"0":[1.0,3.1,4.5]}}' + # The MAP value arrives as a decimal.Decimal on every pandas version, and + # pandas 3 serializes Decimal to a JSON string where pandas 2 emitted a number. + map_value = '"1"' if pandas_major_version >= 3 else "1.0" + assert pdf.to_json() == ( + '{"MAP":{"0":[["k1",' + map_value + "]]}," + '"OBJ":{"0":{"A":"foo","b":0.05}},' + '"ARR":{"0":[1.0,3.1,4.5]}}' ) else: assert ( diff --git a/tests/integ/test_df_to_pandas.py b/tests/integ/test_df_to_pandas.py index 862af71b90..881bfad663 100644 --- a/tests/integ/test_df_to_pandas.py +++ b/tests/integ/test_df_to_pandas.py @@ -53,6 +53,9 @@ ) from tests.utils import IS_IN_STORED_PROC, Utils +# Not imported from _internal.utils: stored procs run the server's released snowpark. +pandas_major_version = int(pd.__version__.split(".")[0]) + def test_to_pandas_new_df_from_range(session): # Single column @@ -104,9 +107,10 @@ def test_to_pandas_cast_integer(session, to_pandas_api, local_testing_mode): assert ( str(pandas_df.dtypes.iloc[4]) == "int64" ) # When limits are not explicitly defined, rely on metadata information from GS. - assert ( - str(pandas_df.dtypes.iloc[5]) == "object" - ) # No cast so it's a string. dtype is "object". + assert str(pandas_df.dtypes.iloc[5]) in ( + "object", + "str", + ) # No cast so it's a string: object on pandas 2, str on pandas 3. assert ( str(pandas_df.dtypes.iloc[6]) == "float64" ) # A 20-digit number is over int64 max. Convert to float64 in pandas. @@ -125,7 +129,6 @@ def test_to_pandas_cast_integer(session, to_pandas_api, local_testing_mode): # Starting from pyarrow 13, pyarrow no longer coerces non-nanosecond to nanosecond for pandas >=2.0 # https://arrow.apache.org/release/13.0.0.html and https://github.com/apache/arrow/issues/33321 pyarrow_major_version = int(pa.__version__.split(".")[0]) - pandas_major_version = int(pd.__version__.split(".")[0]) expected_dtype = ( "datetime64[s]" if pyarrow_major_version >= 13 and pandas_major_version >= 2 @@ -134,7 +137,11 @@ def test_to_pandas_cast_integer(session, to_pandas_api, local_testing_mode): assert str(timestamp_pandas_df.dtypes.iloc[0]) == expected_dtype else: # TODO: mock the non-nanosecond unit pyarrow+pandas behavior in local test - assert str(timestamp_pandas_df.dtypes.iloc[0]) == "datetime64[ns]" + # The mock layer follows pandas' own default resolution: ns on pandas 2, us on 3. + assert str(timestamp_pandas_df.dtypes.iloc[0]) in ( + "datetime64[ns]", + "datetime64[us]", + ) def test_to_pandas_precision_for_number_38_0(session): @@ -270,7 +277,21 @@ def test_to_pandas_batches(session, local_testing_mode): @pytest.mark.skipif( IS_IN_STORED_PROC, reason="SNOW-1362480, backend optimization in different reg env" ) -def test_df_to_pandas_df(session): +def test_df_to_pandas_df(session, local_testing_mode): + # The live path builds the frame through pyarrow, which pins a TIMESTAMP_NTZ + # column to datetime64[ns]. The local testing emulator returns pandas' own + # default resolution instead, and that default moved from ns to us in pandas + # 3, so only the live expectation can be pinned. + timestamp_dtype = None if local_testing_mode else "datetime64[ns]" + # From pandas 3 on, pyarrow maps an Arrow string column onto the dedicated + # str dtype, so an all-NULL column arrives as str rather than object on the + # live path. The local testing emulator never goes through pyarrow. + null_dtype = ( + object + if local_testing_mode or pandas_major_version < 3 + else pd.Series([None], dtype="str").dtype + ) + df = session.create_dataframe( [ [ @@ -319,7 +340,8 @@ def test_df_to_pandas_df(session): minute=12, second=12, ) - ] + ], + dtype=timestamp_dtype, ), } ) @@ -376,7 +398,7 @@ def test_df_to_pandas_df(session): pandas_df = pd.DataFrame( { - "A": pd.Series(["[\n 1,\n 2,\n 3,\n 4\n]"], dtype=object), + "A": pd.Series(["[\n 1,\n 2,\n 3,\n 4\n]"]), "B": pd.Series([b"123"], dtype=object), "C": pd.Series([True], dtype=bool), "D": pd.Series( @@ -392,17 +414,17 @@ def test_df_to_pandas_df(session): "J": pd.Series( [100], dtype=np.int64 ), # in reg env, there can be backend optimization resulting in np.int8 - "K": pd.Series([None], dtype=object), + "K": pd.Series([None], dtype=null_dtype), "L": pd.Series( [100], dtype=np.int64 ), # in reg env, there can be backend optimization resulting in np.int8 - "M": pd.Series(["abc"], dtype=object), + "M": pd.Series(["abc"]), "N": pd.Series( - [datetime.datetime(2023, 10, 30, 12, 12, 12)], dtype="datetime64[ns]" + [datetime.datetime(2023, 10, 30, 12, 12, 12)], dtype=timestamp_dtype ), "O": pd.Series([datetime.time(12, 12, 12)], dtype=object), - "P": pd.Series(['{\n "a": "b"\n}'], dtype=object), - "Q": pd.Series(['{\n "a": "b"\n}'], dtype=object), + "P": pd.Series(['{\n "a": "b"\n}']), + "Q": pd.Series(['{\n "a": "b"\n}']), } ) assert_frame_equal(df.to_pandas(), pandas_df) diff --git a/tests/integ/test_udf.py b/tests/integ/test_udf.py index 7acee3c1c4..737e98421d 100644 --- a/tests/integ/test_udf.py +++ b/tests/integ/test_udf.py @@ -2374,7 +2374,7 @@ def return_type_in_dataframe(x): ("", ""), ("float64",), ), - (StringType, [["1"]], ("",), ("object",)), + (StringType, [["1"]], ("",), ("str", "object")), ( BooleanType, [[True]], @@ -2397,7 +2397,7 @@ def return_type_in_dataframe(x): ("",), ("object",), ), - (ArrayType, [[[1]]], ("",), ("object",)), + (ArrayType, [[[1]]], ("",), ("str", "object")), ( TimeType, [[datetime.time(1, 1, 1)]], @@ -2407,12 +2407,15 @@ def return_type_in_dataframe(x): ( TimestampType, [[datetime.datetime(2016, 3, 13, 5, tzinfo=datetime.timezone.utc)]], - ("",), + ( + "", + "", + ), ("datetime64[ns]",), ), - (GeographyType, [["POINT(30 10)"]], ("",), ("object",)), - (GeometryType, [["POINT(30 10)"]], ("",), ("object",)), - (MapType, [[{1: 2}]], ("",), ("object",)), + (GeographyType, [["POINT(30 10)"]], ("",), ("str", "object")), + (GeometryType, [["POINT(30 10)"]], ("",), ("str", "object")), + (MapType, [[{1: 2}]], ("",), ("str", "object")), ], ) def test_pandas_udf_return_types(session, _type, data, expected_types, expected_dtypes): @@ -2429,7 +2432,7 @@ def test_pandas_udf_return_types(session, _type, data, expected_types, expected_ immutable=True, ) result_df = df.select(series_udf("a")).to_pandas() - result_val = result_df.iloc[0][0] + result_val = result_df.iloc[0, 0] if _type in (ArrayType, MapType, GeographyType, GeometryType): # TODO: SNOW-573478 result_val = json.loads(result_val) @@ -2437,8 +2440,8 @@ def test_pandas_udf_return_types(session, _type, data, expected_types, expected_ str(type(result_val)) in expected_types ), f"returned type is {type(result_val)} instead of {expected_types}" assert ( - result_df.dtypes[0] in expected_dtypes - ), f"returned dtype is {result_df.dtypes[0]} instead of {expected_dtypes}" + result_df.dtypes.iloc[0] in expected_dtypes + ), f"returned dtype is {result_df.dtypes.iloc[0]} instead of {expected_dtypes}" @pytest.mark.skipif( @@ -2481,13 +2484,14 @@ def test_pandas_udf_return_variant(session): input_types=[PandasSeriesType(VariantType())], ) temp = df.select(series_udf("a")).to_pandas() - assert ( - temp.dtypes[0] == object - ), f"returned dtype is {temp.dtypes[0]} instead of object" + assert temp.dtypes.iloc[0] in ( + "str", + "object", + ), f"returned dtype is {temp.dtypes.iloc[0]} instead of str or object" for i, row in temp.iterrows(): assert isinstance( - row[0], expected_types[i] - ), f"returned type is {type(row[0])} instead of {expected_types[i]}" + row.iloc[0], expected_types[i] + ), f"returned type is {type(row.iloc[0])} instead of {expected_types[i]}" @pytest.mark.skipif( diff --git a/tests/mock/test_functions.py b/tests/mock/test_functions.py index 84409240be..3675c9a3be 100644 --- a/tests/mock/test_functions.py +++ b/tests/mock/test_functions.py @@ -27,7 +27,10 @@ dense_rank, desc, get, + greatest, + initcap, is_null, + least, lit, max, min, @@ -36,6 +39,7 @@ sum, to_char, to_date, + udf, ) from snowflake.snowpark.mock._functions import MockedFunctionRegistry, patch from snowflake.snowpark.mock._snowflake_data_type import ColumnEmulator, ColumnType @@ -68,6 +72,29 @@ def test_col(session): assert origin_df.select(col("o")).collect() == [Row(True), Row(False), Row(None)] +def test_initcap_greatest_least_keep_null(session): + df = session.create_dataframe( + [["hello", "a", "b"], [None, None, "b"]], + schema=["name", "a", "b"], + ) + initcap_out = df.select(initcap(col("name")).alias("R")) + assert initcap_out.collect() == [Row("Hello"), Row(None)] + assert initcap_out.filter(col("R").is_null()).count() == 1 + + greatest_out = df.select(greatest(col("a"), col("b")).alias("R")) + assert greatest_out.collect() == [Row("b"), Row(None)] + assert greatest_out.filter(col("R").is_null()).count() == 1 + + least_out = df.select(least(col("a"), col("b")).alias("R")) + assert least_out.collect() == [Row("a"), Row(None)] + assert least_out.filter(col("R").is_null()).count() == 1 + + delim_df = session.create_dataframe([["hello-world"], [None]], schema=["name"]) + delim_out = delim_df.select(initcap(col("name"), lit("-")).alias("R")) + assert delim_out.collect() == [Row("Hello-World"), Row(None)] + assert delim_out.filter(col("R").is_null()).count() == 1 + + def test_max(session): origin_df: DataFrame = session.create_dataframe( [ @@ -884,3 +911,21 @@ def test_save_as_table_column_order_name_array_type(session): assert len(result) == 1 assert result[0]["ID"] == 1 assert result[0]["TAGS"] is None + + +def test_concat_null_does_not_reach_strict_udf(session): + # SNOW-3923354: pandas 3 re-inference leaked the NULL into the UDF as nan. + seen = [] + + @udf(strict=True, return_type=StringType(), input_types=[StringType()]) + def probe(s): + seen.append(s) + return "CALLED" + + df = session.create_dataframe([["a"], [None], ["c"]], schema=["v"]) + assert df.select(probe(concat(col("v"), col("v")))).collect() == [ + Row("CALLED"), + Row(None), + Row("CALLED"), + ] + assert sorted(seen) == ["aa", "cc"] diff --git a/tests/unit/test_code_generation.py b/tests/unit/test_code_generation.py index 52c7e54769..93eaaed5c6 100644 --- a/tests/unit/test_code_generation.py +++ b/tests/unit/test_code_generation.py @@ -592,7 +592,7 @@ def func(): assert ( generate_source_code(func) - == """\ + == f"""\ # The following comment contains the source code generated by snowpark-python for explanatory purposes. # import io # import pandas @@ -601,7 +601,7 @@ def func(): # import pandas.io.sql # import random as rd # from datetime import date as d -# from pandas.io.sql import read_sql +# from {read_sql.__module__} import read_sql # def func(): # pandas.io # io diff --git a/tests/unit/test_session.py b/tests/unit/test_session.py index a11c557ea7..60e84c7ff7 100644 --- a/tests/unit/test_session.py +++ b/tests/unit/test_session.py @@ -1,6 +1,7 @@ # # Copyright (c) 2012-2025 Snowflake Computing Inc. All rights reserved. # +import datetime import json import logging import os @@ -22,6 +23,13 @@ except ImportError: is_pandas_available = False +try: + import pyarrow as pa + + is_pyarrow_available = True +except ImportError: + is_pyarrow_available = False + from snowflake.snowpark import Session from snowflake.snowpark._internal.server_connection import ServerConnection from snowflake.snowpark._internal.utils import parse_table_name @@ -372,6 +380,86 @@ def test_resolve_packages_optional_artifact_repository(mock_server_connection): ] +@pytest.mark.skipif(not is_pyarrow_available, reason="requires pyarrow") +def test_normalize_arrow_duration_to_ns(): + one_second = 1_000_000_000 + table = pa.table( + { + "s": pa.array([1], type=pa.duration("s")), + "ms": pa.array([1_000], type=pa.duration("ms")), + "us": pa.array([1_000_000], type=pa.duration("us")), + "ns": pa.array([one_second], type=pa.duration("ns")), + "id": pa.array([1], type=pa.int64()), + } + ) + out = Session._normalize_arrow_duration_to_ns(table) + assert out.schema.field("id").type == pa.int64() + for name in ("s", "ms", "us", "ns"): + assert out.schema.field(name).type == pa.duration("ns") + assert out.column(name).cast(pa.int64())[0].as_py() == one_second + + +@pytest.mark.skipif(not is_pandas_available, reason="requires pandas for write_pandas") +@pytest.mark.parametrize("unit", ["s", "ms", "us", "ns"]) +def test_write_pandas_normalizes_timedelta_to_ns(mock_server_connection, unit): + session = Session(mock_server_connection) + one_day = datetime.timedelta(days=1) + df = pandas.DataFrame( + { + "D": pandas.Series([one_day]).astype(f"timedelta64[{unit}]"), + "N": [1], + } + ) + captured = {} + + def fake_write_pandas(_conn, frame, *args, **kwargs): + captured["frame"] = frame + return True, 1, 1, [] + + with mock.patch.object( + snowflake.snowpark.session, "write_pandas", side_effect=fake_write_pandas + ): + session.write_pandas(df, table_name="t") + + frame = captured["frame"] + assert str(frame["D"].dtype) == "timedelta64[ns]" + # Snowflake stores the raw tick count with no unit, so ns is the contract. + assert frame["D"].to_numpy().astype("int64")[0] == 86_400_000_000_000 + assert str(frame["N"].dtype) == "int64" + assert str(df["D"].dtype) == f"timedelta64[{unit}]" + + +@pytest.mark.skipif(not is_pyarrow_available, reason="requires pyarrow") +@pytest.mark.parametrize( + "unit,ticks", + [("s", 1), ("ms", 1_000), ("us", 1_000_000), ("ns", 1_000_000_000)], +) +def test_write_arrow_normalizes_duration_to_ns(mock_server_connection, unit, ticks): + session = Session(mock_server_connection) + table = pa.table( + { + "D": pa.array([ticks], type=pa.duration(unit)), + "N": pa.array([1], type=pa.int64()), + } + ) + captured = {} + + def fake_write_arrow(*args, **kwargs): + captured["table"] = kwargs["table"] + return True, 1, 1, [] + + with mock.patch.object( + snowflake.snowpark.session, "write_arrow", side_effect=fake_write_arrow + ): + session.write_arrow(table, table_name="t") + + out = captured["table"] + assert out.schema.field("D").type == pa.duration("ns") + assert out.column("D").cast(pa.int64())[0].as_py() == 1_000_000_000 + assert out.schema.field("N").type == pa.int64() + assert table.schema.field("D").type == pa.duration(unit) + + @pytest.mark.skipif(not is_pandas_available, reason="requires pandas for write_pandas") def test_write_pandas_wrong_table_type(mock_server_connection): session = Session(mock_server_connection) diff --git a/tox.ini b/tox.ini index c77f9710f0..073a992c77 100644 --- a/tox.ini +++ b/tox.ini @@ -57,6 +57,8 @@ deps = {env:SNOWFLAKE_PYTEST_MODIN_DEPS} {env:SNOWFLAKE_PYTEST_MODIN_PIN} {env:SNOWFLAKE_PYTEST_PANDAS_DEPS} + pandas2: pandas>=2.1.2,<3 + pandas3: pandas>=3,<4 install_command = bash ./scripts/tox_install_cmd.sh {opts} {packages} setenv = COVERAGE_FILE = {env:COVERAGE_FILE:{toxworkdir}/.coverage.{envname}} @@ -129,6 +131,10 @@ passenv = GITHUB_ENV SNOWPARK_PYTHON_API_TEST_BUCKET_PATH SNOWPARK_PYTHON_API_S3_STORAGE_INTEGRATION +commands_pre = + pandas2: python -c "import pandas; assert pandas.__version__.startswith('2.'), f'wanted pandas 2, got {pandas.__version__}'" + pandas3: python -c "import pandas; assert pandas.__version__.startswith('3.'), f'wanted pandas 3, got {pandas.__version__}'" + local: python -c "import pandas, pyarrow; print('pandas', pandas.__version__, 'pyarrow', pyarrow.__version__)" commands = notudf: {env:SNOWFLAKE_PYTEST_CMD} -m "{env:SNOWFLAKE_TEST_TYPE} and not udf" {posargs:} {env:RERUN_FLAGS} src/snowflake/snowpark tests udf: {env:SNOWFLAKE_PYTEST_CMD} -m "{env:SNOWFLAKE_TEST_TYPE} or udf" {posargs:} {env:RERUN_FLAGS} src/snowflake/snowpark tests