dan's homework

Author: rok

python/pyarrow-stubs/pyarrow/_parquet.pyi

@@ -467,6 +467,7 @@ class ParquetWriter(_Weakrefable):
         sorting_columns: tuple[SortingColumn, ...] | None = None,
         store_decimal_as_integer: bool = False,
         write_time_adjusted_to_utc: bool = False,
+        max_rows_per_page: int | None = None,
     ): ...
     def close(self) -> None: ...
     def write_table(self, table: Table, row_group_size: int | None = None) -> None: ...

python/pyarrow-stubs/pyarrow/_parquet_encryption.pyi

@@ -16,9 +16,11 @@
 # under the License.
 
 import datetime as dt
+import pathlib
 
 from collections.abc import Callable
 
+from pyarrow._fs import FileSystem
 from ._parquet import FileDecryptionProperties, FileEncryptionProperties
 from .lib import _Weakrefable
 
@@ -73,7 +75,7 @@ class KmsConnectionConfig(_Weakrefable):
 
 class KmsClient(_Weakrefable):
     def wrap_key(self, key_bytes: bytes, master_key_identifier: str) -> str: ...
-    def unwrap_key(self, wrapped_key: str, master_key_identifier: str) -> str: ...
+    def unwrap_key(self, wrapped_key: str, master_key_identifier: str) -> bytes: ...
 
 
 class CryptoFactory(_Weakrefable):
@@ -93,3 +95,47 @@ class CryptoFactory(_Weakrefable):
     ) -> FileDecryptionProperties: ...
     def remove_cache_entries_for_token(self, access_token: str) -> None: ...
     def remove_cache_entries_for_all_tokens(self) -> None: ...
+    def rotate_master_keys(
+        self,
+        kms_connection_config: KmsConnectionConfig,
+        parquet_file_path: str | pathlib.Path,
+        filesystem: FileSystem | None = None,
+        double_wrapping: bool = True,
+        cache_lifetime_seconds: int | float = 600,
+    ) -> None: ...
+
+
+class KeyMaterial(_Weakrefable):
+    @property
+    def is_footer_key(self) -> bool: ...
+    @property
+    def is_double_wrapped(self) -> bool: ...
+    @property
+    def master_key_id(self) -> str: ...
+    @property
+    def wrapped_dek(self) -> str: ...
+    @property
+    def kek_id(self) -> str: ...
+    @property
+    def wrapped_kek(self) -> str: ...
+    @property
+    def kms_instance_id(self) -> str: ...
+    @property
+    def kms_instance_url(self) -> str: ...
+    @staticmethod
+    def wrap(key_material: KeyMaterial) -> KeyMaterial: ...
+    @staticmethod
+    def parse(key_material_string: str) -> KeyMaterial: ...
+
+
+
+class FileSystemKeyMaterialStore(_Weakrefable):
+    def get_key_material(self, key_id: str) -> KeyMaterial: ...
+    def get_key_id_set(self) -> list[str]: ...
+    @classmethod
+    def for_file(
+            cls,
+            parquet_file_path: str | pathlib.Path, /,
+            filesystem: FileSystem | None = None
+    ) -> FileSystemKeyMaterialStore:
+        ...

python/pyarrow-stubs/pyarrow/compute.pyi

@@ -318,6 +318,7 @@ def first(
     memory_pool: lib.MemoryPool | None = None,
 ) -> _ScalarT: ...
 
+last = _clone_signature(first)
 
 def first_last(
     array: lib.Array[Any] | lib.ChunkedArray[Any] | list[Any],
@@ -546,6 +547,7 @@ def exp(
     _FloatArrayT | lib.DoubleArray | _FloatScalarT | lib.DoubleScalar | Expression): ...
 
 
+expm1 = _clone_signature(exp)
 multiply = _clone_signature(add)
 multiply_checked = _clone_signature(add)
 
@@ -741,10 +743,12 @@ logb_checked = _clone_signature(logb)
 acos = _clone_signature(ln)
 acos_checked = _clone_signature(ln)
 acosh = _clone_signature(ln)
+acosh_checked = _clone_signature(ln)
 asin = _clone_signature(ln)
 asin_checked = _clone_signature(ln)
 asinh = _clone_signature(ln)
 atan = _clone_signature(ln)
+atanh_checked = _clone_signature(ln)
 atanh = _clone_signature(ln)
 cos = _clone_signature(ln)
 cos_checked = _clone_signature(ln)
@@ -1171,6 +1175,13 @@ def index_in(
     memory_pool: lib.MemoryPool | None = None,
 ) -> lib.Int32Scalar | lib.Int32Array | Expression: ...
 
+def index_in_meta_binary(
+    values: lib.Scalar | lib.Array | lib.ChunkedArray | Expression,
+    value_set: lib.Array | lib.ChunkedArray | Expression,
+    /,
+    *,
+    memory_pool: lib.MemoryPool | None = None,
+) -> lib.Int32Scalar | lib.Int32Array | Expression: ...
 
 def is_in(
     values: lib.Scalar | lib.Array | lib.ChunkedArray | Expression,
@@ -1183,6 +1194,7 @@ def is_in(
 ) -> lib.BooleanScalar | lib.BooleanArray | Expression: ...
 
 
+is_in_meta_binary = _clone_signature(index_in_meta_binary)
 match_like = _clone_signature(ends_with)
 match_substring = _clone_signature(ends_with)
 match_substring_regex = _clone_signature(ends_with)

python/pyarrow-stubs/pyarrow/parquet/core.pyi

@@ -198,6 +198,7 @@ class ParquetWriter:
         write_page_checksum: bool = False,
         sorting_columns: Sequence[SortingColumn] | None = None,
         store_decimal_as_integer: bool = False,
+        max_rows_per_page: int | None = None,
         **options,
     ) -> None: ...
     def __enter__(self) -> Self: ...

python/pyarrow-stubs/pyarrow/parquet/encryption.pyi

@@ -19,6 +19,7 @@ from pyarrow._parquet_encryption import (
     CryptoFactory,
     DecryptionConfiguration,
     EncryptionConfiguration,
+    FileSystemKeyMaterialStore,
     KmsClient,
     KmsConnectionConfig,
 )
@@ -27,6 +28,7 @@ __all__ = [
     "CryptoFactory",
     "DecryptionConfiguration",
     "EncryptionConfiguration",
+    "FileSystemKeyMaterialStore",
     "KmsClient",
     "KmsConnectionConfig",
 ]

python/pyarrow/parquet/core.py

@@ -1908,7 +1908,9 @@ def read_table(source, *, columns=None, use_threads=True,
 
         filesystem, path = _resolve_filesystem_and_path(source, filesystem)
         if filesystem is not None:
-            if not filesystem.get_file_info(path).is_file:
+            file_info = filesystem.get_file_info(path)
+            assert isinstance(file_info, FileInfo)
+            if not file_info.is_file:
                 raise ValueError(
                     "the 'source' argument should be "
                     "an existing parquet file and not a directory "

python/pyarrow/parquet/encryption.py

@@ -20,4 +20,5 @@
                                          EncryptionConfiguration,
                                          DecryptionConfiguration,
                                          KmsConnectionConfig,
-                                         KmsClient)
+                                         KmsClient,
+                                         FileSystemKeyMaterialStore)

python/pyarrow/tests/parquet/test_encryption.py

@@ -553,7 +553,7 @@ def test_encrypted_parquet_write_read_external(tempdir, data_table,
     result_table = read_encrypted_parquet(
         path, decryption_config, kms_connection_config, crypto_factory,
         internal_key_material=False)
-    store = pa._parquet_encryption.FileSystemKeyMaterialStore.for_file(path)
+    store = pe.FileSystemKeyMaterialStore.for_file(path)
 
     assert len(key_ids := store.get_key_id_set()) == (
         len(external_encryption_config.column_keys[COL_KEY_NAME]) + 1)

python/pyarrow/tests/test_flight.py

@@ -266,6 +266,7 @@ class EchoTableStreamFlightServer(EchoFlightServer):
     """An echo server that streams the whole table."""
 
     def do_get(self, context, ticket):
+        assert self.last_message is not None
         return flight.GeneratorStream(
             self.last_message.schema,
             [self.last_message])
@@ -283,6 +284,7 @@ class EchoRecordBatchReaderStreamFlightServer(EchoFlightServer):
     """An echo server that streams the whole table as a RecordBatchReader."""
 
     def do_get(self, context, ticket):
+        assert self.last_message is not None
         return flight.GeneratorStream(
             self.last_message.schema,
             [self.last_message.to_reader()])