fix: PyArrow import broken in Lambda (#2337)

Author: LeonLuttenberger

awswrangler/s3/_read_orc.py

@@ -17,7 +17,6 @@
 import pandas as pd
 import pyarrow as pa
 import pyarrow.dataset
-import pyarrow.orc
 from typing_extensions import Literal
 
 from awswrangler import _data_types, _utils, exceptions
@@ -40,16 +39,19 @@
 
 if TYPE_CHECKING:
     from mypy_boto3_s3 import S3Client
+    from pyarrow.orc import ORCFile
 
 FULL_READ_S3_BLOCK_SIZE = 20_971_520  # 20 MB (20 * 2**20)
 METADATA_READ_S3_BLOCK_SIZE = 131_072  # 128 KB (128 * 2**10)
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
 
-def _pyarrow_orc_file_wrapper(source: Any) -> pyarrow.orc.ORCFile:
+def _pyarrow_orc_file_wrapper(source: Any) -> "ORCFile":
+    from pyarrow.orc import ORCFile
+
     try:
-        return pyarrow.orc.ORCFile(source=source)
+        return ORCFile(source=source)
     except pyarrow.ArrowInvalid as ex:
         if str(ex) == "ORC file size is 0 bytes":
             _logger.warning("Ignoring empty file...")
@@ -74,7 +76,7 @@ def _read_orc_metadata_file(
         s3_block_size=METADATA_READ_S3_BLOCK_SIZE,
         s3_additional_kwargs=s3_additional_kwargs,
     ) as f:
-        orc_file: Optional[pyarrow.orc.ORCFile] = _pyarrow_orc_file_wrapper(source=f)
+        orc_file: Optional["ORCFile"] = _pyarrow_orc_file_wrapper(source=f)
         if orc_file:
             return orc_file.schema
         return None
@@ -118,7 +120,7 @@ def _read_orc_file(
         s3_additional_kwargs=s3_additional_kwargs,
         s3_client=s3_client,
     ) as f:
-        orc_file: Optional[pyarrow.orc.ORCFile] = _pyarrow_orc_file_wrapper(
+        orc_file: Optional["ORCFile"] = _pyarrow_orc_file_wrapper(
             source=f,
         )
         if orc_file is None:

awswrangler/s3/_write_orc.py

@@ -8,7 +8,6 @@
 import boto3
 import pandas as pd
 import pyarrow as pa
-import pyarrow.orc
 
 from awswrangler import _utils, catalog, exceptions, typing
 from awswrangler._arrow import _df_to_table
@@ -33,6 +32,7 @@
 
 if TYPE_CHECKING:
     from mypy_boto3_s3 import S3Client
+    from pyarrow.orc import ORCWriter
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -54,8 +54,10 @@ def _new_writer(
     s3_client: "S3Client",
     s3_additional_kwargs: Optional[Dict[str, str]],
     use_threads: Union[bool, int],
-) -> Iterator[pyarrow.orc.ORCWriter]:
-    writer: Optional[pyarrow.orc.ORCWriter] = None
+) -> Iterator["ORCWriter"]:
+    from pyarrow.orc import ORCWriter
+
+    writer: Optional["ORCWriter"] = None
     if not pyarrow_additional_kwargs:
         pyarrow_additional_kwargs = {}
 
@@ -67,7 +69,7 @@ def _new_writer(
         s3_client=s3_client,
     ) as f:
         try:
-            writer = pyarrow.orc.ORCWriter(
+            writer = ORCWriter(
                 where=f,
                 compression="uncompressed" if compression is None else compression,
                 **pyarrow_additional_kwargs,