/ mlflow / utils / spark_utils.py
spark_utils.py
 1  def is_spark_connect_mode():
 2      try:
 3          from pyspark.sql.utils import is_remote
 4      except ImportError:
 5          return False
 6      return is_remote()
 7  
 8  
 9  def get_spark_dataframe_type():
10      if is_spark_connect_mode():
11          from pyspark.sql.connect.dataframe import DataFrame as SparkDataFrame
12      else:
13          from pyspark.sql import DataFrame as SparkDataFrame
14  
15      return SparkDataFrame