From 3dd4fa9fdde7f7150907baa966355534546d1976 Mon Sep 17 00:00:00 2001 From: Leon Luttenberger Date: Wed, 14 Jun 2023 15:47:57 -0600 Subject: [PATCH] fix: Import for arrow_orc_datasource --- .../distributed/ray/datasources/arrow_orc_datasource.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/awswrangler/distributed/ray/datasources/arrow_orc_datasource.py b/awswrangler/distributed/ray/datasources/arrow_orc_datasource.py index a6fd868a4..739d966f4 100644 --- a/awswrangler/distributed/ray/datasources/arrow_orc_datasource.py +++ b/awswrangler/distributed/ray/datasources/arrow_orc_datasource.py @@ -2,7 +2,6 @@ from typing import Any, Dict, List, Optional import pyarrow as pa -from pyarrow import orc from ray.data.block import BlockAccessor from awswrangler._arrow import _add_table_partitions, _df_to_table @@ -21,6 +20,8 @@ def _read_file( # type: ignore[override] path_root: str, **reader_args: Any, ) -> pa.Table: + from pyarrow import orc + columns: Optional[List[str]] = reader_args.get("columns", None) table: pa.Table = orc.read_table(f, columns=columns) @@ -48,6 +49,8 @@ def _write_block( # type: ignore[override] pandas_kwargs: Optional[Dict[str, Any]], **writer_args: Any, ) -> None: + from pyarrow import orc + schema: Optional[pa.schema] = writer_args.get("schema", None) dtype: Optional[Dict[str, str]] = writer_args.get("dtype", None) index: bool = writer_args.get("index", False)