checkout merged dataframe.py

Author: kosiew

python/datafusion/dataframe.py

@@ -22,16 +22,13 @@
 from __future__ import annotations
 
 import warnings
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from typing import (
     TYPE_CHECKING,
     Any,
     AsyncIterator,
-    Iterable,
     Iterator,
     Literal,
-    Optional,
-    Union,
     overload,
 )
 
@@ -59,7 +56,7 @@
 
 if TYPE_CHECKING:
     import pathlib
-    from typing import Callable
+    from collections.abc import Callable
 
     import pandas as pd
     import polars as pl
@@ -82,7 +79,7 @@ class Compression(Enum):
     LZ4 = "lz4"
     # lzo is not implemented yet
     # https://github.com/apache/arrow-rs/issues/6970
-    # LZO = "lzo"
+    # LZO = "lzo"  # noqa: ERA001
     ZSTD = "zstd"
     LZ4_RAW = "lz4_raw"
 
@@ -109,7 +106,7 @@ def from_str(cls: type[Compression], value: str) -> Compression:
                 """
             raise ValueError(error_msg) from err
 
-    def get_default_level(self) -> Optional[int]:
+    def get_default_level(self) -> int | None:
         """Get the default compression level for the compression type.
 
         Returns:
@@ -142,24 +139,24 @@ def __init__(
         write_batch_size: int = 1024,
         writer_version: str = "1.0",
         skip_arrow_metadata: bool = False,
-        compression: Optional[str] = "zstd(3)",
-        compression_level: Optional[int] = None,
-        dictionary_enabled: Optional[bool] = True,
+        compression: str | None = "zstd(3)",
+        compression_level: int | None = None,
+        dictionary_enabled: bool | None = True,
         dictionary_page_size_limit: int = 1024 * 1024,
-        statistics_enabled: Optional[str] = "page",
+        statistics_enabled: str | None = "page",
         max_row_group_size: int = 1024 * 1024,
         created_by: str = "datafusion-python",
-        column_index_truncate_length: Optional[int] = 64,
-        statistics_truncate_length: Optional[int] = None,
+        column_index_truncate_length: int | None = 64,
+        statistics_truncate_length: int | None = None,
         data_page_row_count_limit: int = 20_000,
-        encoding: Optional[str] = None,
+        encoding: str | None = None,
         bloom_filter_on_write: bool = False,
-        bloom_filter_fpp: Optional[float] = None,
-        bloom_filter_ndv: Optional[int] = None,
+        bloom_filter_fpp: float | None = None,
+        bloom_filter_ndv: int | None = None,
         allow_single_file_parallelism: bool = True,
         maximum_parallel_row_group_writers: int = 1,
         maximum_buffered_record_batches_per_stream: int = 2,
-        column_specific_options: Optional[dict[str, ParquetColumnOptions]] = None,
+        column_specific_options: dict[str, ParquetColumnOptions] | None = None,
     ) -> None:
         """Initialize the ParquetWriterOptions.
 
@@ -264,13 +261,13 @@ class ParquetColumnOptions:
 
     def __init__(
         self,
-        encoding: Optional[str] = None,
-        dictionary_enabled: Optional[bool] = None,
-        compression: Optional[str] = None,
-        statistics_enabled: Optional[str] = None,
-        bloom_filter_enabled: Optional[bool] = None,
-        bloom_filter_fpp: Optional[float] = None,
-        bloom_filter_ndv: Optional[int] = None,
+        encoding: str | None = None,
+        dictionary_enabled: bool | None = None,
+        compression: str | None = None,
+        statistics_enabled: str | None = None,
+        bloom_filter_enabled: bool | None = None,
+        bloom_filter_fpp: float | None = None,
+        bloom_filter_ndv: int | None = None,
     ) -> None:
         """Initialize the ParquetColumnOptions.
 
@@ -836,7 +833,7 @@ def join(
         # of a keyword argument.
         if (
             isinstance(on, tuple)
-            and len(on) == 2
+            and len(on) == 2  # noqa: PLR2004
             and isinstance(on[0], list)
             and isinstance(on[1], list)
         ):
@@ -1068,7 +1065,7 @@ def write_parquet(
     def write_parquet(
         self,
         path: str | pathlib.Path,
-        compression: Union[str, Compression, ParquetWriterOptions] = Compression.ZSTD,
+        compression: str | Compression | ParquetWriterOptions = Compression.ZSTD,
         compression_level: int | None = None,
         write_options: DataFrameWriteOptions | None = None,
     ) -> None: