SNOW-2334855: mock_get returns the datatype of the container, not of the returned value

[pedro-villanueva-bcom]

Please answer these questions before submitting your issue. Thanks!

1. What version of Python are you using?

Python 3.12.8

2. What are the Snowpark Python and pandas versions in the environment?

pandas==2.3.1
snowflake-snowpark-python==1.35.0

3. What did you do?

from snowflake.snowpark import functions as F
from snowflake.snowpark import types as T
from snowflake.snowpark import Session
import pandas as pd

from snowflake.snowpark.functions import get
from snowflake.snowpark.mock import ColumnEmulator, ColumnType, patch

FIX = False

if FIX:

    @patch(get)
    def mock_get(column_expression: ColumnEmulator, value_expression: ColumnEmulator) -> ColumnEmulator:
        """Correctly mock get returning the right type"""

        def get(obj, key):
            try:
                if isinstance(obj, list) and key < len(obj):
                    return obj[key]
                elif isinstance(obj, dict):
                    return obj.get(key, None)
                else:
                    return None
            except KeyError:
                return None

        # pandas.Series.combine does not work here because it will not allow Nones in int columns
        result = []
        for exp, k in zip(column_expression, value_expression):
            result.append(get(exp, k))

        if column_expression.sf_type.datatype == T.ArrayType():
            return_type = column_expression.sf_type.datatype.element_type
        elif column_expression.sf_type.datatype == T.MapType():
            return_type = column_expression.sf_type.datatype.valueType
        else:
            raise ValueError(f"Only arrays and maps are supported, not {column_expression.sf_type.datatype}")

        return ColumnEmulator(result, sf_type=ColumnType(return_type, True), dtype=object)


spark = Session.builder.config("local_testing", True).getOrCreate()


def test_get():
    """Test the mocked get function"""
    data = pd.DataFrame(
        [
            [1, ["a", "b", "c"], 1],
            [2, ["d", "e"], 1],
            [3, ["f", "g", "h"], 1],
        ],
        columns=["ID", "COLUMN", "POS"],
    )
    df_array = spark.createDataFrame(data).select(
        "ID", F.col("COLUMN").cast("array<string>").alias("COLUMN"), "POS"
    )

    df_array.select(F.get(F.col("COLUMN"), F.col("POS")).alias("VALUE")).print_schema()

    data = pd.DataFrame(
        [
            [1, {"a": 1, "b": 3, "c": 5}, "a"],
            [2, {"d": 8, "e": 2}, "d"],
            [3, {"f": 5, "g": 7, "h": 1}, "f"],
        ],
        columns=["ID", "COLUMN", "KEY"],
    )
    df_map = spark.createDataFrame(data).select(
        "ID", F.col("COLUMN").cast("map<string,int>").alias("COLUMN"), "KEY"
    )
    df_map.select(F.get(F.col("COLUMN"), F.col("KEY")).alias("VALUE")).print_schema()


test_get()

4. What did you expect to see?

With FIX=False, I get

root
 |-- "VALUE": ArrayType (nullable = True)
 |   |-- element: StringType()
root
 |-- "VALUE": MapType (nullable = True)
 |   |-- key: StringType()
 |   |-- value: StringType()

Which is wrong. The get function extracts elements from arrays and maps, so the returned column should have the data type of the contained object not the container.

With FIX=True. The expected and logical output is this:

root
 |-- "VALUE": StringType() (nullable = True)
root
 |-- "VALUE": StringType() (nullable = True)

I've fixed the mocked get function by returning the contained type:

if column_expression.sf_type.datatype == T.ArrayType():
    return_type = column_expression.sf_type.datatype.element_type
elif column_expression.sf_type.datatype == T.MapType():
    return_type = column_expression.sf_type.datatype.valueType
else:
    raise ValueError(f"Only arrays and maps are supported, not {column_expression.sf_type.datatype}")

return ColumnEmulator(result, sf_type=ColumnType(return_type, True), dtype=object)

I can open a PR to change the code if you think this is the right approach.

[sfc-gh-sghosh]

Hi @pedro-villanueva-bcom ,

Thanks for raising the issue.
While trying to reproduce the issue with local session, whether the FIX is True or False, in both case, the output is String data type, where as for live session (for both case True or False), the output is having variant data type.

local:
--- Testing with an ArrayType column ---
Schema of the array result DataFrame:
root
|-- "ID": LongType() (nullable = True)
|-- "VALUE": StringType() (nullable = True)

Array result:

|"ID" |"VALUE" |

|1 |b |
|2 |e |
|3 |g |

--- Testing with a MapType column ---
Schema of the map result DataFrame:
root
|-- "ID": LongType() (nullable = True)
|-- "VALUE": StringType() (nullable = True)

Map result:

|"ID" |"VALUE" |

|1 |1 |
|2 |8 |
|3 |5 |

Live session:

--- Testing with an ArrayType column ---
Schema of the array result DataFrame:
root
|-- "ID": LongType() (nullable = True)
|-- "VALUE": VariantType() (nullable = True)

Array result:

|"ID" |"VALUE" |

|1 |"b" |
|2 |"e" |
|3 |"g" |

--- Testing with a MapType column ---
Schema of the map result DataFrame:
root
|-- "ID": LongType() (nullable = True)
|-- "VALUE": VariantType() (nullable = True)

Map result:

|"ID" |"VALUE" |

|1 |1 |
|2 |8 |
|3 |5 |

[pedro-villanueva-bcom]

Hello,
What python, pandas, and snowpark versions are you using?

[pedro-villanueva-bcom]

Also I'm a bit confused. This fix doesn't change the live results at all, so it's fine. VariantType is expected I believe.
For the local results that you show, did you set FIX to true or false? The results that you have look like FIX=True. The key is that the output schema is correctly set to a string, not an array<string> or map<string>.