2222from __future__ import annotations
2323
2424import warnings
25- from collections .abc import Iterable , Sequence
25+ from collections .abc import Sequence
2626from typing import (
2727 TYPE_CHECKING ,
2828 Any ,
2929 AsyncIterator ,
30+ Iterable ,
3031 Iterator ,
3132 Literal ,
33+ Optional ,
34+ Union ,
3235 overload ,
3336)
3437
5659
5760if TYPE_CHECKING :
5861 import pathlib
59- from collections . abc import Callable
62+ from typing import Callable
6063
6164 import pandas as pd
6265 import polars as pl
@@ -79,7 +82,7 @@ class Compression(Enum):
7982 LZ4 = "lz4"
8083 # lzo is not implemented yet
8184 # https://github.com/apache/arrow-rs/issues/6970
82- # LZO = "lzo" # noqa: ERA001
85+ # LZO = "lzo"
8386 ZSTD = "zstd"
8487 LZ4_RAW = "lz4_raw"
8588
@@ -106,7 +109,7 @@ def from_str(cls: type[Compression], value: str) -> Compression:
106109 """
107110 raise ValueError (error_msg ) from err
108111
109- def get_default_level (self ) -> int | None :
112+ def get_default_level (self ) -> Optional [ int ] :
110113 """Get the default compression level for the compression type.
111114
112115 Returns:
@@ -139,24 +142,24 @@ def __init__(
139142 write_batch_size : int = 1024 ,
140143 writer_version : str = "1.0" ,
141144 skip_arrow_metadata : bool = False ,
142- compression : str | None = "zstd(3)" ,
143- compression_level : int | None = None ,
144- dictionary_enabled : bool | None = True ,
145+ compression : Optional [ str ] = "zstd(3)" ,
146+ compression_level : Optional [ int ] = None ,
147+ dictionary_enabled : Optional [ bool ] = True ,
145148 dictionary_page_size_limit : int = 1024 * 1024 ,
146- statistics_enabled : str | None = "page" ,
149+ statistics_enabled : Optional [ str ] = "page" ,
147150 max_row_group_size : int = 1024 * 1024 ,
148151 created_by : str = "datafusion-python" ,
149- column_index_truncate_length : int | None = 64 ,
150- statistics_truncate_length : int | None = None ,
152+ column_index_truncate_length : Optional [ int ] = 64 ,
153+ statistics_truncate_length : Optional [ int ] = None ,
151154 data_page_row_count_limit : int = 20_000 ,
152- encoding : str | None = None ,
155+ encoding : Optional [ str ] = None ,
153156 bloom_filter_on_write : bool = False ,
154- bloom_filter_fpp : float | None = None ,
155- bloom_filter_ndv : int | None = None ,
157+ bloom_filter_fpp : Optional [ float ] = None ,
158+ bloom_filter_ndv : Optional [ int ] = None ,
156159 allow_single_file_parallelism : bool = True ,
157160 maximum_parallel_row_group_writers : int = 1 ,
158161 maximum_buffered_record_batches_per_stream : int = 2 ,
159- column_specific_options : dict [str , ParquetColumnOptions ] | None = None ,
162+ column_specific_options : Optional [ dict [str , ParquetColumnOptions ]] = None ,
160163 ) -> None :
161164 """Initialize the ParquetWriterOptions.
162165
@@ -261,13 +264,13 @@ class ParquetColumnOptions:
261264
262265 def __init__ (
263266 self ,
264- encoding : str | None = None ,
265- dictionary_enabled : bool | None = None ,
266- compression : str | None = None ,
267- statistics_enabled : str | None = None ,
268- bloom_filter_enabled : bool | None = None ,
269- bloom_filter_fpp : float | None = None ,
270- bloom_filter_ndv : int | None = None ,
267+ encoding : Optional [ str ] = None ,
268+ dictionary_enabled : Optional [ bool ] = None ,
269+ compression : Optional [ str ] = None ,
270+ statistics_enabled : Optional [ str ] = None ,
271+ bloom_filter_enabled : Optional [ bool ] = None ,
272+ bloom_filter_fpp : Optional [ float ] = None ,
273+ bloom_filter_ndv : Optional [ int ] = None ,
271274 ) -> None :
272275 """Initialize the ParquetColumnOptions.
273276
@@ -833,7 +836,7 @@ def join(
833836 # of a keyword argument.
834837 if (
835838 isinstance (on , tuple )
836- and len (on ) == 2 # noqa: PLR2004
839+ and len (on ) == 2
837840 and isinstance (on [0 ], list )
838841 and isinstance (on [1 ], list )
839842 ):
@@ -1065,7 +1068,7 @@ def write_parquet(
10651068 def write_parquet (
10661069 self ,
10671070 path : str | pathlib .Path ,
1068- compression : str | Compression | ParquetWriterOptions = Compression .ZSTD ,
1071+ compression : Union [ str , Compression , ParquetWriterOptions ] = Compression .ZSTD ,
10691072 compression_level : int | None = None ,
10701073 write_options : DataFrameWriteOptions | None = None ,
10711074 ) -> None :
0 commit comments