Remove rdd

quinn/dataframe_helpers.py

@@ -27,7 +27,7 @@ def column_to_list(df: DataFrame, col_name: str) -> list[Any]:
 
     # sparksession from df is not available in older versions of pyspark
     if sys.modules["pyspark"].__version__ < "3.3.0":
-        return df.select(col_name).rdd.flatMap(lambda x: x).collect()
+        return [row[0] for row in df.select(col_name).collect()]
 
     spark_config = df.sparkSession.sparkContext.getConf().getAll()
 
@@ -40,7 +40,7 @@ def column_to_list(df: DataFrame, col_name: str) -> list[Any]:
     if pyarrow_valid and pandas_valid:
         return df.select(col_name).toPandas()[col_name].tolist()
 
-    return df.select(col_name).rdd.flatMap(lambda x: x).collect()
+    return [row[0] for row in df.select(col_name).collect()]
 
 
 def two_columns_to_dictionary(

quinn/transformations.py

@@ -252,7 +252,7 @@ def fix_nullability(field: StructField, result_dict: dict) -> None:
         spark = SparkSession.getActiveSession()
         spark = spark if spark is not None else SparkSession.builder.getOrCreate()
 
-    return spark.createDataFrame(output.rdd, output.schema)
+    return output
 
 
 def flatten_struct(df: DataFrame, col_name: str, separator: str = ":") -> DataFrame: