star

venv/lib/python3.9/site-packages/mlflow/utils/spark_utils.py

@@ -0,0 +1,15 @@
+def is_spark_connect_mode():
+    try:
+        from pyspark.sql.utils import is_remote
+    except ImportError:
+        return False
+    return is_remote()
+
+
+def get_spark_dataframe_type():
+    if is_spark_connect_mode():
+        from pyspark.sql.connect.dataframe import DataFrame as SparkDataFrame
+    else:
+        from pyspark.sql import DataFrame as SparkDataFrame
+
+    return SparkDataFrame