Skip to content

Commit 28e0ee8

Browse files
committed
use .api
1 parent 52ff113 commit 28e0ee8

File tree

1 file changed

+13
-17
lines changed

1 file changed

+13
-17
lines changed

pandas/io/parquet.py

+13-17
Original file line numberDiff line numberDiff line change
@@ -24,25 +24,24 @@ def __init__(self):
2424
# we need to import on first use
2525

2626
try:
27-
import pyarrow # noqa
27+
import pyarrow
28+
from pyarrow import parquet
2829
except ImportError:
2930
raise ImportError("pyarrow is required for parquet support\n\n"
3031
"you can install via conda\n"
3132
"conda install pyarrow -c conda-forge\n"
3233
"\nor via pip\n"
3334
"pip install pyarrow\n")
35+
self.api = pyarrow
36+
self.parquet = parquet
3437

3538
def write(self, df, path, compression=None, **kwargs):
36-
import pyarrow
37-
from pyarrow import parquet as pq
38-
39-
table = pyarrow.Table.from_pandas(df)
40-
pq.write_table(table, path,
41-
compression=compression, **kwargs)
39+
table = self.api.Table.from_pandas(df)
40+
self.parquet.write_table(
41+
table, path, compression=compression, **kwargs)
4242

4343
def read(self, path):
44-
import pyarrow
45-
return pyarrow.parquet.read_table(path).to_pandas()
44+
return self.parquet.read_table(path).to_pandas()
4645

4746

4847
class FastParquetImpl(object):
@@ -52,28 +51,25 @@ def __init__(self):
5251
# we need to import on first use
5352

5453
try:
55-
import fastparquet # noqa
54+
import fastparquet
5655
except ImportError:
5756
raise ImportError("fastparquet is required for parquet support\n\n"
5857
"you can install via conda\n"
5958
"conda install fastparquet -c conda-forge\n"
6059
"\nor via pip\n"
6160
"pip install fastparquet")
61+
self.api = fastparquet
6262

6363
def write(self, df, path, compression=None, **kwargs):
64-
import fastparquet
65-
6664
# thriftpy/protocol/compact.py:339:
6765
# DeprecationWarning: tostring() is deprecated.
6866
# Use tobytes() instead.
6967
with catch_warnings(record=True):
70-
fastparquet.write(path, df,
71-
compression=compression, **kwargs)
68+
self.api.write(path, df,
69+
compression=compression, **kwargs)
7270

7371
def read(self, path):
74-
import fastparquet
75-
pf = fastparquet.ParquetFile(path)
76-
return pf.to_pandas()
72+
return self.api.ParquetFile(path).to_pandas()
7773

7874

7975
def to_parquet(df, path, engine, compression=None, **kwargs):

0 commit comments

Comments
 (0)