SNOW-3153913: patching split() in local testing

[serban-dobroiu]

Hi, trying to patch and use split() for testing some dataframe transformations but I'm not sure I'm doing it right:

@patch(split)
def mock_split(column: ColumnEmulator, pattern: ColumnEmulator) -> ColumnEmulator:
    ret_column = ColumnEmulator(data=[split(row, pattern=pattern) for row in column])
    ret_column.sf_type = ColumnType(ArrayType(StringType()), nullable=True)
    return ret_column

When I try to test a transformation function calling split():

def col_to_array(
    df: DataFrame,
    cols: Iterable[str],
    delimiter: str = "|",
    ignore_missing: bool = False,
) -> DataFrame:
....
for _col in cols:
        transformation = when(col(_col).is_not_null(), split(col(_col), lit(delimiter)))
....

it is throwing the following type error:

error_message = "Error executing mocked function 'split'. See error traceback for detailed information."```

Can you please help me understand what is happening here?

thank you!

[sfc-gh-sghosh]

Hello @serban-dobroiu ,

Thanks for raising the issue. Could you please share the import libraries as well? Which version of Snowpark Python are you using, which verison of python ?

[serban-dobroiu]

@sfc-gh-sghosh sorry about that

• here are the imports

import pytest
from unittest.mock import MagicMock, create_autospec
from snowflake.snowpark import Session
from snowflake.snowpark.dataframe import DataFrame
from snowflake.snowpark.types import (
    ArrayType,
    StructType,
    StructField,
    StringType,
    IntegerType,
    FloatType,
)
from snowflake.snowpark.functions import split
from snowflake.snowpark.mock import patch, ColumnEmulator, ColumnType
from snowflake.snowpark.testing import assert_dataframe_equal
from utilities.transformation_functions import col_to_array


# add parent to Path
import sys
from pathlib import Path

and Python 3.11.5 and Snowflake Snowpark Python 1.44.0

Thank you!

[sfc-gh-sghosh]

Hi @serban-dobroiu ,

I tested with snowpark 1.44.0 and python 3.12 and used str.split() instead of Snowpark's split()

I am not getting any error, although there is little difference with the split function for local session compare to regular session, but there is no error

Here 

import os
from typing import Iterable
from snowflake.snowpark import Session
from snowflake.snowpark.dataframe import DataFrame
from snowflake.snowpark.types import (
    ArrayType,
    StructType,
    StructField,
    StringType,
)
from snowflake.snowpark.functions import split, col, lit, when

USE_LOCAL_TESTING = True

if USE_LOCAL_TESTING:
    from snowflake.snowpark.mock import patch, ColumnEmulator, ColumnType
    session = Session.builder.config("local_testing", True).create()
else:
    import snowflake.connector
    session = Session.builder.configs({
        "connection": snowflake.connector.connect(
            connection_name=os.getenv("SNOWFLAKE_CONNECTION_NAME") or "default"
        )
    }).create()

if USE_LOCAL_TESTING:
    @patch(split)
    def mock_split(column: ColumnEmulator, pattern: ColumnEmulator) -> ColumnEmulator:
        delimiter = str(pattern.iloc[0]) if len(pattern) > 0 else "|"
        ret_column = ColumnEmulator(data=[str(row).split(delimiter) for row in column])
        ret_column.sf_type = ColumnType(ArrayType(StringType()), nullable=True)
        return ret_column

def col_to_array(
    df: DataFrame,
    cols: Iterable[str],
    delimiter: str = "|",
    ignore_missing: bool = False,
) -> DataFrame:
    for _col in cols:
        transformation = when(
            col(_col).is_not_null(), split(col(_col), lit(delimiter))
        )
        df = df.with_column(_col, transformation)
    return df

test_data = session.create_dataframe(
    [
        ("apple|banana|cherry",),
        ("dog|cat",),
        (None,),
        ("single",),
    ],
    schema=StructType([StructField("FRUITS", StringType())]),
)

print("Original data:")
test_data.show()

print("\nAfter col_to_array:")
result = col_to_array(test_data, cols=["FRUITS"], delimiter="|")
result.show()
print("SUCCESS")

Output:
**Regular session:**

Original data:
|"FRUITS" |
|apple|banana|cherry |
|dog|cat |
|NULL |
|single |
After col_to_array:
|"FRUITS" |
|[ |
| "apple", |
| "banana", |
| "cherry" |
|] |
|[ |
| "dog", |
| "cat" |
|] |
|NULL |
|[ |
| "single" |
|] |
SUCCESS

**local session**

Original data:
|"FRUITS" |
|apple|banana|cherry |
|dog|cat |
|NULL |
|single |
After col_to_array:
|"FRUITS" |
|['apple', 'banana', 'cherry'] |
|['dog', 'cat'] |
|NULL |
|['single'] |

[serban-dobroiu]

@sfc-gh-sghosh thank you! I will give that a go shortly.

In the meantime, this is something that's not clear to me; functionally you are testing the Python string split() function not the Snowpark DataFrame split() - but my feeling is this is not quite the same? I'm not challenging anything, it's just that I am trying to understand if these are equivalent approaches and to what extent?

thank you!