move imports out of functions

Author: kosiew

python/datafusion/dataframe.py

@@ -38,14 +38,15 @@
 from datafusion.plan import ExecutionPlan, LogicalPlan
 from datafusion.record_batch import RecordBatchStream
 
+import pyarrow as pa
+from datafusion import functions as f
 
 if TYPE_CHECKING:
     import pathlib
     from typing import Callable, Sequence
 
     import pandas as pd
     import polars as pl
-    import pyarrow as pa
 
 from enum import Enum
 
@@ -874,8 +875,6 @@ def fill_null(self, value: Any, subset: list[str] | None = None) -> "DataFrame":
             - For columns where casting fails, the original column is kept unchanged
             - For columns not in subset, the original column is kept unchanged
         """
-        import pyarrow as pa
-        from datafusion import functions as f
 
         # Get columns to process
         if subset is None:
@@ -910,38 +909,39 @@ def fill_null(self, value: Any, subset: list[str] | None = None) -> "DataFrame":
                 exprs.append(f.col(col_name))
 
         return self.select(*exprs)
-    
-    def fill_nan(self, value: float | int, subset: list[str] | None = None) -> "DataFrame":
+
+    def fill_nan(
+        self, value: float | int, subset: list[str] | None = None
+    ) -> "DataFrame":
         """Fill NaN values in specified numeric columns with a value.
-        
+
         Args:
             value: Numeric value to replace NaN values with
             subset: Optional list of column names to fill. If None, fills all numeric columns.
-        
+
         Returns:
             DataFrame with NaN values replaced in numeric columns
-        
+
         Examples:
             >>> df = df.fill_nan(0)  # Fill all NaNs with 0 in numeric columns
             >>> df = df.fill_nan(99.9, subset=["price", "score"])  # Fill specific columns
-        
+
         Notes:
             - Only fills NaN values in numeric columns (float32, float64)
             - Non-numeric columns are kept unchanged
             - For columns not in subset, the original column is kept unchanged
             - Value must be numeric (int or float)
         """
-        import pyarrow as pa
-        from datafusion import functions as f
-        
+
         if not isinstance(value, (int, float)):
             raise ValueError("Value must be numeric (int or float)")
 
         # Get columns to process
         if subset is None:
             # Only get numeric columns if no subset specified
             subset = [
-                field.name for field in self.schema() 
+                field.name
+                for field in self.schema()
                 if pa.types.is_floating(field.type)
             ]
         else:
@@ -962,5 +962,5 @@ def fill_nan(self, value: float | int, subset: list[str] | None = None) -> "Data
             else:
                 # Keep columns not in subset unchanged
                 exprs.append(f.col(col_name))
-                
+
         return self.select(*exprs)

python/tests/test_dataframe.py

@@ -1264,7 +1264,7 @@ def test_fill_null(df):
     )
     with pytest.raises(ValueError, match="Column 'f' not found in DataFrame"):
         df_with_nulls.fill_null("missing", subset=["e", "f"])
-    
+
     def test_fill_nan(df):
         # Test filling NaNs with integer value
         df_with_nans = df.with_column("d", literal(float("nan")).cast(pa.float64()))
@@ -1311,4 +1311,4 @@ def test_fill_nan(df):
         df_filled = df_with_nans.fill_nan(99.9, subset=["e"])
         result = df_filled.to_pydict()
         assert result["d"] == [float("nan"), float("nan"), float("nan")]
-        assert result["e"] == [99.9, 99.9, 99.9]
+        assert result["e"] == [99.9, 99.9, 99.9]