snowpark-connect 0.22.1__py3-none-any.whl → 0.23.0__py3-none-any.whl

snowflake/snowpark_connect/relation/read/utils.py

@@ -73,12 +73,13 @@ def rename_columns_as_snowflake_standard(
         return df, []
 
     new_columns = make_column_names_snowpark_compatible(df.columns, plan_id)
-    return (
-        df.select(
-            *(df.col(orig).alias(alias) for orig, alias in zip(df.columns, new_columns))
-        ),
-        new_columns,
-    )
+    result = df.toDF(*new_columns)
+    if result._select_statement is not None:
+        # do not allow snowpark to flatten the to_df result
+        # TODO: remove after SNOW-2203706 is fixed
+        result._select_statement.flatten_disabled = True
+
+    return (result, new_columns)
 
 
 class Connection(Protocol):

snowflake/snowpark_connect/relation/utils.py

@@ -6,12 +6,28 @@ import random
 import re
 import string
 import time
-from typing import Sequence
+from collections.abc import Callable
+from typing import AbstractSet, List, Optional, Sequence
 
 import pyspark.sql.connect.proto.relations_pb2 as relation_proto
 
 import snowflake.snowpark.functions as snowpark_fn
 from snowflake import snowpark
+from snowflake.snowpark._internal.analyzer.expression import (
+    COLUMN_DEPENDENCY_ALL,
+    COLUMN_DEPENDENCY_DOLLAR,
+    Expression,
+    FunctionExpression,
+    derive_dependent_columns,
+)
+from snowflake.snowpark._internal.analyzer.select_statement import (
+    SEQUENCE_DEPENDENT_DATA_GENERATION,
+    ColumnChangeState,
+    ColumnStateDict,
+    SelectStatement,
+)
+from snowflake.snowpark._internal.analyzer.unary_expression import Alias
+from snowflake.snowpark._internal.analyzer.window_expression import WindowExpression
 from snowflake.snowpark.types import (
     BinaryType,
     BooleanType,
@@ -223,3 +239,156 @@ def snowpark_functions_col(name: str, column_map: ColumnNameMap) -> snowpark.Col
     """
     is_qualified_name = name not in column_map.get_snowpark_columns()
     return snowpark_fn.col(name, _is_qualified_name=is_qualified_name)
+
+
+def can_sort_be_flattened(
+    select_statement: Optional[SelectStatement], *sort_expressions: Optional[Expression]
+) -> bool:
+    """
+    Checks if the given SelectStatement can be "flattened" when sorting with regard to the given sort expressions.
+    Flattening means that the given SelectStatement can be enhanced and reused instead of being treated
+    as a subquery in the FROM clause after a "sort" or "filter" operation. Flattening allows accessing dropped columns
+    for sort and filter expressions.
+    """
+    if not select_statement or select_statement.flatten_disabled:
+        return False
+
+    # In some cases, flattening sort can lead to leaving the "order by" clause in a subquery,
+    # which can cause incorrect ordering. We want to avoid flattening sort when all its dependent columns
+    # are available in the current projection.
+    dependent_columns_in_sort = derive_dependent_columns(*sort_expressions)
+    columns_in_projection = _get_columns_in_projection(select_statement.projection)
+    if len(dependent_columns_in_sort - columns_in_projection) == 0:
+        return False
+
+    return _can_clause_dependent_columns_flatten(
+        dependent_columns_in_sort, select_statement.column_states
+    ) and not _has_data_generator_exp(select_statement.projection)
+
+
+def can_filter_be_flattened(
+    select_statement: Optional[SelectStatement], condition: Expression
+) -> bool:
+    """
+    Checks if the given SelectStatement can be "flattened" when filtering with regard to the given condition.
+    Flattening means that the given SelectStatement can be enhanced and reused instead of being treated
+    as a subquery in the FROM clause after a "sort" or "filter" operation. Flattening allows accessing dropped columns
+    for sort and filter expressions.
+    """
+    if not select_statement or select_statement.flatten_disabled:
+        return False
+
+    return all(
+        [
+            _can_clause_dependent_columns_flatten(
+                derive_dependent_columns(condition), select_statement.column_states
+            ),
+            not _has_data_generator_or_window_exp(select_statement.projection),
+            select_statement.order_by is None,
+            select_statement.limit_ is None,
+        ]
+    )
+
+
+def _get_columns_in_projection(
+    projection: Optional[List[Expression]],
+) -> AbstractSet[str]:
+    if projection is None:
+        return set()
+
+    columns = set()
+    for expression in projection:
+        if hasattr(expression, "name") and expression.name:
+            columns.add(expression.name)
+        elif hasattr(expression, "children"):
+            columns.update(_get_columns_in_projection(expression.children))
+
+    return columns
+
+
+def _is_self_alias(expression):
+    """
+    Check if the expression is a self-alias, meaning it has an alias that is the same as its name.
+    A self-alias can be flattened, even if Snowpark treats it as a CHANGED_EXP.
+    """
+    if not isinstance(expression, Alias):
+        return False
+
+    first_child_with_name = expression.child
+    while (
+        first_child_with_name
+        and hasattr(first_child_with_name, "child")
+        and not hasattr(first_child_with_name, "name")
+    ):
+        first_child_with_name = first_child_with_name.child
+
+    return (
+        first_child_with_name
+        and (first_child_with_name, "name")
+        and first_child_with_name.name == expression.name
+    )
+
+
+def _can_clause_dependent_columns_flatten(
+    dependent_columns: Optional[AbstractSet[str]],
+    subquery_column_states: ColumnStateDict,
+) -> bool:
+    if dependent_columns == COLUMN_DEPENDENCY_DOLLAR:
+        return False
+    elif (
+        subquery_column_states.has_changed_columns
+        or subquery_column_states.has_new_columns
+    ):
+        if dependent_columns == COLUMN_DEPENDENCY_ALL:
+            return False
+
+        assert dependent_columns is not None
+        for dc in dependent_columns:
+            dc_state = subquery_column_states.get(dc)
+            if dc_state:
+                if (
+                    dc_state.change_state == ColumnChangeState.CHANGED_EXP
+                    and not _is_self_alias(dc_state.expression)
+                ):
+                    return False
+    return True
+
+
+def _has_data_generator_exp(expressions: List[Expression]) -> bool:
+    return _has_expression(expressions, [_is_generator_expression])
+
+
+def _has_data_generator_or_window_exp(expressions: List[Expression]) -> bool:
+    return _has_expression(
+        expressions, [_is_generator_expression, _is_window_expression]
+    )
+
+
+def _has_expression(
+    expressions: Optional[List[Expression]], checks: List[Callable[[Expression], bool]]
+) -> bool:
+    if expressions is None:
+        return False
+
+    for exp in expressions:
+        if not exp:
+            continue
+
+        if any([check(exp) for check in checks]):
+            return True
+
+        if _has_expression(exp.children, checks):
+            return True
+
+    return False
+
+
+def _is_window_expression(exp: Expression) -> bool:
+    return isinstance(exp, WindowExpression)
+
+
+def _is_generator_expression(exp: Expression) -> bool:
+    # https://docs.snowflake.com/en/sql-reference/functions-data-generation
+    return isinstance(exp, FunctionExpression) and (
+        exp.is_data_generator or exp.name.lower() in SEQUENCE_DEPENDENT_DATA_GENERATION
+    )

snowflake/snowpark_connect/relation/write/map_write.py

@@ -234,6 +234,10 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         column_order=_column_order_for_write,
                     )
                 case "append":
+                    if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                        raise AnalysisException(
+                            f"Table {snowpark_table_name} is not an iceberg table"
+                        )
                     _validate_schema_and_get_writer(
                         input_df, "append", snowpark_table_name
                     ).saveAsTable(
@@ -260,7 +264,10 @@ def map_write(request: proto_base.ExecutePlanRequest):
                         )
                 case "overwrite":
                     if check_snowflake_table_existence(snowpark_table_name, session):
-                        session.sql(f"DELETE FROM {snowpark_table_name}").collect()
+                        if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                            raise AnalysisException(
+                                f"Table {snowpark_table_name} is not an iceberg table"
+                            )
                     else:
                         create_iceberg_table(
                             snowpark_table_name=snowpark_table_name,
@@ -269,10 +276,10 @@ def map_write(request: proto_base.ExecutePlanRequest):
                             snowpark_session=session,
                         )
                     _validate_schema_and_get_writer(
-                        input_df, "append", snowpark_table_name
+                        input_df, "truncate", snowpark_table_name
                     ).saveAsTable(
                         table_name=snowpark_table_name,
-                        mode="append",
+                        mode="truncate",
                         column_order=_column_order_for_write,
                     )
                 case _:
@@ -286,13 +293,46 @@ def map_write(request: proto_base.ExecutePlanRequest):
                 write_op.table.save_method
                 == commands_proto.WriteOperation.SaveTable.TableSaveMethod.TABLE_SAVE_METHOD_SAVE_AS_TABLE
             ):
-                _validate_schema_and_get_writer(
-                    input_df, write_mode, snowpark_table_name
-                ).saveAsTable(
-                    table_name=snowpark_table_name,
-                    mode=write_mode,
-                    column_order=_column_order_for_write,
-                )
+                match write_mode:
+                    case "overwrite":
+                        if check_snowflake_table_existence(
+                            snowpark_table_name, session
+                        ):
+                            if (
+                                check_table_type(snowpark_table_name, session)
+                                != "TABLE"
+                            ):
+                                raise AnalysisException(
+                                    f"Table {snowpark_table_name} is not a FDN table"
+                                )
+                            write_mode = "truncate"
+                        _validate_schema_and_get_writer(
+                            input_df, write_mode, snowpark_table_name
+                        ).saveAsTable(
+                            table_name=snowpark_table_name,
+                            mode=write_mode,
+                            column_order=_column_order_for_write,
+                        )
+                    case "append":
+                        if check_table_type(snowpark_table_name, session) != "TABLE":
+                            raise AnalysisException(
+                                f"Table {snowpark_table_name} is not a FDN table"
+                            )
+                        _validate_schema_and_get_writer(
+                            input_df, write_mode, snowpark_table_name
+                        ).saveAsTable(
+                            table_name=snowpark_table_name,
+                            mode=write_mode,
+                            column_order=_column_order_for_write,
+                        )
+                    case _:
+                        _validate_schema_and_get_writer(
+                            input_df, write_mode, snowpark_table_name
+                        ).saveAsTable(
+                            table_name=snowpark_table_name,
+                            mode=write_mode,
+                            column_order=_column_order_for_write,
+                        )
             elif (
                 write_op.table.save_method
                 == commands_proto.WriteOperation.SaveTable.TableSaveMethod.TABLE_SAVE_METHOD_INSERT_INTO
@@ -312,21 +352,6 @@ def map_write(request: proto_base.ExecutePlanRequest):
 
 def map_write_v2(request: proto_base.ExecutePlanRequest):
     write_op = request.plan.command.write_operation_v2
-    match write_op.mode:
-        case commands_proto.WriteOperationV2.MODE_APPEND:
-            write_mode = "append"
-        case commands_proto.WriteOperationV2.MODE_CREATE:
-            write_mode = "errorifexists"
-        case commands_proto.WriteOperationV2.MODE_OVERWRITE:
-            write_mode = "overwrite"
-        case commands_proto.WriteOperationV2.MODE_REPLACE:
-            write_mode = "overwrite"
-        case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
-            write_mode = "overwrite"
-        case _:
-            raise SnowparkConnectNotImplementedError(
-                f"Write operation {write_op.mode} not implemented."
-            )
 
     snowpark_table_name = _spark_to_snowflake(write_op.table_name)
     result = map_relation(write_op.input)
@@ -338,55 +363,176 @@ def map_write_v2(request: proto_base.ExecutePlanRequest):
             "Write operation V2 only support table writing now"
         )
 
-    # For OVERWRITE and APPEND modes, check if table exists first - Spark requires table to exist for these operations
-    if write_op.mode in (
-        commands_proto.WriteOperationV2.MODE_OVERWRITE,
-        commands_proto.WriteOperationV2.MODE_APPEND,
-    ):
-        if not check_snowflake_table_existence(snowpark_table_name, session):
-            raise AnalysisException(
-                f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found. "
-                f"Verify the spelling and correctness of the schema and catalog.\n"
-            )
-
     if write_op.provider.lower() == "iceberg":
-        if write_mode == "overwrite" and check_snowflake_table_existence(
-            snowpark_table_name, session
-        ):
-            session.sql(f"DELETE FROM {snowpark_table_name}").collect()
-            write_mode = "append"
-
-        if write_mode in (
-            "errorifexists",
-            "overwrite",
-        ):
-            create_iceberg_table(
-                snowpark_table_name=snowpark_table_name,
-                location=write_op.table_properties.get("location"),
-                schema=input_df.schema,
-                snowpark_session=session,
-            )
-        _validate_schema_and_get_writer(
-            input_df, write_mode, snowpark_table_name
-        ).saveAsTable(
-            table_name=snowpark_table_name,
-            mode="append",
-            column_order=_column_order_for_write,
-        )
+        match write_op.mode:
+            case commands_proto.WriteOperationV2.MODE_CREATE:
+                if check_snowflake_table_existence(snowpark_table_name, session):
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} already exists"
+                    )
+                create_iceberg_table(
+                    snowpark_table_name=snowpark_table_name,
+                    location=write_op.table_properties.get("location"),
+                    schema=input_df.schema,
+                    snowpark_session=session,
+                )
+                _validate_schema_and_get_writer(
+                    input_df, "append", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_APPEND:
+                if not check_snowflake_table_existence(snowpark_table_name, session):
+                    raise AnalysisException(
+                        f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
+                    )
+                if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} is not an iceberg table"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "append", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
+                # TODO: handle the filter condition for MODE_OVERWRITE
+                if check_snowflake_table_existence(snowpark_table_name, session):
+                    if check_table_type(snowpark_table_name, session) != "ICEBERG":
+                        raise AnalysisException(
+                            f"Table {snowpark_table_name} is not an iceberg table"
+                        )
+                else:
+                    raise AnalysisException(
+                        f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "truncate", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="truncate",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_REPLACE:
+                if check_snowflake_table_existence(snowpark_table_name, session):
+                    create_iceberg_table(
+                        snowpark_table_name=snowpark_table_name,
+                        location=write_op.table_properties.get("location"),
+                        schema=input_df.schema,
+                        snowpark_session=session,
+                        mode="replace",
+                    )
+                else:
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} does not exist"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "replace", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
+                create_iceberg_table(
+                    snowpark_table_name=snowpark_table_name,
+                    location=write_op.table_properties.get("location"),
+                    schema=input_df.schema,
+                    snowpark_session=session,
+                    mode="create_or_replace",
+                )
+                _validate_schema_and_get_writer(
+                    input_df, "create_or_replace", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            case _:
+                raise SnowparkConnectNotImplementedError(
+                    f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
+                )
     else:
-        _validate_schema_and_get_writer(
-            input_df, write_mode, snowpark_table_name
-        ).saveAsTable(
-            table_name=snowpark_table_name,
-            mode=write_mode,
-            column_order=_column_order_for_write,
-        )
+        match write_op.mode:
+            case commands_proto.WriteOperationV2.MODE_CREATE:
+                _validate_schema_and_get_writer(
+                    input_df, "errorifexists", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="errorifexists",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_APPEND:
+                if not check_snowflake_table_existence(snowpark_table_name, session):
+                    raise AnalysisException(
+                        f"[TABLE_OR_VIEW_NOT_FOUND] The table or view `{write_op.table_name}` cannot be found."
+                    )
+                if check_table_type(snowpark_table_name, session) != "TABLE":
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} is not a FDN table"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "append", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="append",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_OVERWRITE | commands_proto.WriteOperationV2.MODE_OVERWRITE_PARTITIONS:
+                # TODO: handle the filter condition for MODE_OVERWRITE
+                if check_snowflake_table_existence(snowpark_table_name, session):
+                    if check_table_type(snowpark_table_name, session) != "TABLE":
+                        raise AnalysisException(
+                            f"Table {snowpark_table_name} is not a FDN table"
+                        )
+                else:
+                    raise AnalysisException(
+                        f"[TABLE_OR_VIEW_NOT_FOUND] Table {snowpark_table_name} does not exist"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "truncate", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="truncate",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_REPLACE:
+                if not check_snowflake_table_existence(snowpark_table_name, session):
+                    raise AnalysisException(
+                        f"Table {snowpark_table_name} does not exist"
+                    )
+                _validate_schema_and_get_writer(
+                    input_df, "replace", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="overwrite",
+                    column_order=_column_order_for_write,
+                )
+            case commands_proto.WriteOperationV2.MODE_CREATE_OR_REPLACE:
+                _validate_schema_and_get_writer(
+                    input_df, "create_or_replace", snowpark_table_name
+                ).saveAsTable(
+                    table_name=snowpark_table_name,
+                    mode="overwrite",
+                    column_order=_column_order_for_write,
+                )
+            case _:
+                raise SnowparkConnectNotImplementedError(
+                    f"Write mode {commands_proto.WriteOperationV2.Mode.Name(write_op.mode)} is not supported"
+                )
 
 
 def _validate_schema_and_get_writer(
     input_df: snowpark.DataFrame, write_mode: str, snowpark_table_name: str
 ) -> snowpark.DataFrameWriter:
-    if write_mode == "overwrite":
+    if write_mode is not None and write_mode.lower() in (
+        "replace",
+        "create_or_replace",
+    ):
         return input_df.write
 
     table_schema = None
@@ -528,6 +674,7 @@ def create_iceberg_table(
     location: str,
     schema: StructType,
     snowpark_session: snowpark.Session,
+    mode: str = "create",
 ):
     table_schema = [
         f"{spark_to_sf_single_id(unquote_if_quoted(field.name), is_column = True)} {snowpark_to_iceberg_type(field.datatype)}"
@@ -550,8 +697,20 @@ def create_iceberg_table(
         else f"EXTERNAL_VOLUME = '{config_external_volume}'"
     )
 
+    match mode:
+        case "create":
+            create_sql = "CREATE"
+        case "replace":
+            # There's no replace for iceberg table, so we use create or replace
+            create_sql = "CREATE OR REPLACE"
+        case "create_or_replace":
+            create_sql = "CREATE OR REPLACE"
+        case _:
+            raise SnowparkConnectNotImplementedError(
+                f"Write mode {mode} is not supported for iceberg table"
+            )
     sql = f"""
-        CREATE ICEBERG TABLE {snowpark_table_name} ({",".join(table_schema)})
+        {create_sql} ICEBERG TABLE {snowpark_table_name} ({",".join(table_schema)})
         CATALOG = 'SNOWFLAKE'
         {external_volume}
         {base_location};
@@ -637,3 +796,19 @@ def check_snowflake_table_existence(
         return True
     except Exception:
         return False
+
+
+def check_table_type(
+    snowpark_table_name: str,
+    snowpark_session: snowpark.Session,
+) -> str:
+    # currently we only support iceberg table and FDN table
+    metadata = snowpark_session.sql(
+        f"SHOW TABLES LIKE '{unquote_if_quoted(snowpark_table_name)}';"
+    ).collect()
+    if metadata is None or len(metadata) == 0:
+        raise AnalysisException(f"Table {snowpark_table_name} does not exist")
+    metadata = metadata[0]
+    if metadata.as_dict().get("is_iceberg") == "Y":
+        return "ICEBERG"
+    return "TABLE"

snowflake/snowpark_connect/server.py

@@ -83,6 +83,7 @@ from snowflake.snowpark_connect.utils.context import (
     set_session_id,
     set_spark_version,
 )
+from snowflake.snowpark_connect.utils.env_utils import get_int_from_env
 from snowflake.snowpark_connect.utils.interrupt import (
     interrupt_all_queries,
     interrupt_queries_with_tag,
@@ -700,11 +701,27 @@ def _serve(
             return
 
         server_options = [
-            ("grpc.max_receive_message_length", _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE),
-            ("grpc.max_metadata_size", _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE),
+            (
+                "grpc.max_receive_message_length",
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_MESSAGE_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_MESSAGE_SIZE,
+                ),
+            ),
+            (
+                "grpc.max_metadata_size",
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_METADATA_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE,
+                ),
+            ),
             (
                 "grpc.absolute_max_metadata_size",
-                _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE * 2,
+                get_int_from_env(
+                    "SNOWFLAKE_GRPC_MAX_METADATA_SIZE",
+                    _SPARK_CONNECT_GRPC_MAX_METADATA_SIZE,
+                )
+                * 2,
             ),
         ]
         server = grpc.server(
@@ -812,8 +829,11 @@ class UnixDomainSocketChannelBuilder(ChannelBuilder):
     Spark Connect gRPC channel builder for Unix domain sockets
     """
 
-    def __init__(self, channelOptions: Optional[List[Tuple[str, Any]]] = None) -> None:
-        url: str = get_client_url()
+    def __init__(
+        self, url: str = None, channelOptions: Optional[List[Tuple[str, Any]]] = None
+    ) -> None:
+        if url is None:
+            url = get_client_url()
         if url[:6] != "unix:/" or len(url) < 7:
             raise PySparkValueError(
                 error_class="INVALID_CONNECT_URL",

snowflake/snowpark_connect/type_mapping.py

@@ -59,7 +59,7 @@ def _get_struct_type_class():
 
 
 @cache
-def _get_python_sql_utils_class():
+def get_python_sql_utils_class():
     return jpype.JClass("org.apache.spark.sql.api.python.PythonSQLUtils")
 
 
@@ -70,7 +70,7 @@ def parse_ddl_with_spark_scala(ddl_string: str) -> pyspark.sql.types.DataType:
     This mimics pysparks.ddl parsing logic pyspark.sql.types._py_parse_datatype_string
     """
     struct_type_class = _get_struct_type_class()
-    python_sql_utils = _get_python_sql_utils_class()
+    python_sql_utils = get_python_sql_utils_class()
 
     try:
         # DDL format, "fieldname datatype, fieldname datatype".