Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

SNOW-1865926: Infer schema for StructType columns from nested Rows #2805

Merged
merged 26 commits into from
Jan 7, 2025
Merged

SNOW-1865926: Infer schema for StructType columns from nested Rows #2805

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
d9bf2cb
SNOW-1829870: Allow structured types to be enabled by default
sfc-gh-jrose Dec 5, 2024
ec43e1a
type checking
sfc-gh-jrose Dec 6, 2024
7f3a5fd
lint
sfc-gh-jrose Dec 6, 2024
2e0dce9
Merge branch 'main' into jrose_snow_1829870_structured_by_default
sfc-gh-jrose Dec 16, 2024
ed232de
Move flag to context
sfc-gh-jrose Dec 16, 2024
0dd7b91
typo
sfc-gh-jrose Dec 16, 2024
13c1424
SNOW-1852779 Fix AST encoding for Column `in_`, `asc`, and `desc` (#2…
sfc-gh-vbudati Dec 16, 2024
a787e74
Merge branch 'main' into jrose_snow_1829870_structured_by_default
sfc-gh-jrose Dec 16, 2024
b32806f
merge main and fix test
sfc-gh-jrose Dec 17, 2024
c3db223
make feature flag thread safe
sfc-gh-jrose Dec 17, 2024
1c262d7
typo
sfc-gh-jrose Dec 17, 2024
869931f
Merge branch 'main' into jrose_snow_1829870_structured_by_default
sfc-gh-jrose Dec 17, 2024
0caef58
Fix ast test
sfc-gh-jrose Dec 17, 2024
2380040
move lock
sfc-gh-jrose Dec 18, 2024
995e519
test coverage
sfc-gh-jrose Dec 18, 2024
1b89027
remove context manager
sfc-gh-jrose Dec 18, 2024
4fc61d4
Merge branch 'main' into jrose_snow_1829870_structured_by_default
sfc-gh-jrose Dec 19, 2024
26fd29e
switch to using patch
sfc-gh-jrose Dec 19, 2024
9295e11
move test to other module
sfc-gh-jrose Dec 19, 2024
fcd16d7
Merge branch 'main' into jrose_snow_1829870_structured_by_default
sfc-gh-jrose Dec 19, 2024
77a57a6
fix broken import
sfc-gh-jrose Dec 19, 2024
4769169
another broken import
sfc-gh-jrose Dec 19, 2024
af5af87
another test fix
sfc-gh-jrose Dec 19, 2024
dea741b
Merge branch 'main' into jrose_snow_1829870_structured_by_default
sfc-gh-jrose Dec 20, 2024
ee22980
SNOW-1865926: Infer schema for StructType columns from nested Rows
sfc-gh-jrose Dec 20, 2024
be73744
Merge branch 'main' into jrose_snow_1865926_create_dataframe_default_…
sfc-gh-jrose Jan 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/snowflake/snowpark/_internal/type_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
from snowflake.connector.cursor import ResultMetadata
from snowflake.connector.options import installed_pandas, pandas
from snowflake.snowpark._internal.utils import quote_name
from snowflake.snowpark.row import Row
from snowflake.snowpark.types import (
LTZ,
NTZ,
Expand Down Expand Up @@ -441,6 +442,8 @@ def infer_type(obj: Any) -> DataType:
if key is not None and value is not None:
return MapType(infer_type(key), infer_type(value))
return MapType(NullType(), NullType())
elif isinstance(obj, Row) and context._should_use_structured_type_semantics():
return infer_schema(obj)
Comment on lines +445 to +446
Copy link
Contributor

@sfc-gh-aling sfc-gh-aling Jan 4, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

is there a chance that the given datatype is not StructType while users still input a Row as data? if so what would happen, do we error out?

or Row data always auto inferred as StructType?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm not sure what that scenario would look like. Can you give an example?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

sorry I didn't make it clear, actually I had two questions:

  1. when the infer schema logic would be triggered for Row values -- is it only when schema is not explicitly set?

  2. for my own learning purposes, will following Row input + MapType be a valid input?

    struct = Row(f1="v1", f2=2)
    df = structured_type_session.create_dataframe(
        [
            (struct),
        ],
        schema=[StructureType(MapType(xxx)],
    )

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  1. Yes, the schema is only inferred if no explicitly set.
  2. Today this example would give an error like this:
>>> struct = Row(f1="v1", f2=2)
>>> df = session.create_dataframe([(struct,),], schema=StructType([StructField("test", MapType())]))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "src/snowflake/snowpark/_internal/utils.py", line 960, in func_call_wrapper
    return func(*args, **kwargs)
  File "src/snowflake/snowpark/session.py", line 3318, in create_dataframe
    raise TypeError(
TypeError: Cannot cast <class 'snowflake.snowpark.row.Row'>(Row(f1='v1', f2=2)) to MapType(StringType(), StringType()).

Currently createDataFrame does not know how to handle casting Rows to Maps. You could get it to work by calling Row.as_dict if you wanted it to be a MapType.

After this change it's improved slightly:

# Inferred
>>> df = session.create_dataframe([(struct),])
>>> df.schema
StructType([StructField('F1', StringType(), nullable=False), StructField('F2', LongType(), nullable=False)])
 
# or explicit
>>> df = session.create_dataframe([(struct),], schema=StructType([StructField('F1', StringType(), nullable=False), StructField('F2', LongType(), nullable=False)]))

Without this change the explicit schema still works.

elif isinstance(obj, (list, tuple)):
for v in obj:
if v is not None:
Expand Down
9 changes: 9 additions & 0 deletions src/snowflake/snowpark/session.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,7 @@
import pkg_resources

import snowflake.snowpark._internal.proto.generated.ast_pb2 as proto
import snowflake.snowpark.context as context
from snowflake.connector import ProgrammingError, SnowflakeConnection
from snowflake.connector.options import installed_pandas, pandas
from snowflake.connector.pandas_tools import write_pandas
Expand Down Expand Up @@ -3306,6 +3307,14 @@ def convert_row_to_list(
data_type, (MapType, StructType)
):
converted_row.append(json.dumps(value, cls=PythonObjJSONEncoder))
elif (
isinstance(value, Row)
and isinstance(data_type, StructType)
and context._should_use_structured_type_semantics()
):
converted_row.append(
json.dumps(value.as_dict(), cls=PythonObjJSONEncoder)
)
elif isinstance(data_type, VariantType):
converted_row.append(json.dumps(value, cls=PythonObjJSONEncoder))
elif isinstance(data_type, GeographyType):
Expand Down
41 changes: 41 additions & 0 deletions tests/integ/scala/test_datatype_suite.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -598,6 +598,47 @@ def finish(self) -> dict:
assert MapCollector._return_type == MapType()


@pytest.mark.skipif(
"config.getoption('local_testing_mode', default=False)",
reason="local testing does not fully support structured types yet.",
)
def test_structured_type_infer(structured_type_session, structured_type_support):
if not structured_type_support:
pytest.skip("Test requires structured type support.")

struct = Row(f1="v1", f2=2)
df = structured_type_session.create_dataframe(
[
({"key": "value"}, [1, 2, 3], struct),
],
schema=["map", "array", "obj"],
)

assert df.schema == StructType(
[
StructField(
"MAP",
MapType(StringType(), StringType(), structured=True),
nullable=True,
),
StructField("ARRAY", ArrayType(LongType(), structured=True), nullable=True),
StructField(
"OBJ",
StructType(
[
StructField("f1", StringType(), nullable=True),
StructField("f2", LongType(), nullable=True),
],
structured=True,
),
nullable=True,
),
],
structured=True,
)
df.collect()


@pytest.mark.skipif(
"config.getoption('local_testing_mode', default=False)",
reason="local testing does not fully support structured types yet.",
Expand Down
Loading