[DOP-20162] Change estimate_dataframe_size signature

MobileTeleSystems · Sep 26, 2024 · d4f3a1f · d4f3a1f
1 parent 4188808
commit d4f3a1f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/onetl/_util/spark.py b/onetl/_util/spark.py
@@ -145,9 +145,10 @@ def estimate_dataframe_size(df: DataFrame) -> int:
 
     Using Spark's `SizeEstimator <https://spark.apache.org/docs/3.5.3/api/java/org/apache/spark/util/SizeEstimator.html>`_.
     """
-    spark_session = df._session
+
     try:
-        size_estimator = spark_session._jvm.org.apache.spark.util.SizeEstimator  # type: ignore[union-attr]
+        spark_context = df._sc
+        size_estimator = spark_context._jvm.org.apache.spark.util.SizeEstimator  # type: ignore[union-attr]
         return size_estimator.estimate(df._jdf)
     except Exception:
         # SizeEstimator uses Java reflection which may behave differently in different Java versions,