Use temporary file in load_table_from_dataframe (googleapis#7545)

tswast · web-flow · commit 675dbd88f4dc · 2019-03-25T09:07:18.000-07:00
* Use temporary file in load_table_from_dataframe

This fixes a bug where `load_table_from_dataframe` could not be used
with the `fastparquet` library. It should also use less memory when
uploading large dataframes.

* Add tests using the fastparquet engine.
diff --git a/bigquery/docs/snippets.py b/bigquery/docs/snippets.py
@@ -30,6 +30,10 @@
 import pytest
 import six
 
+try:
+    import fastparquet
+except (ImportError, AttributeError):
+    fastparquet = None
 try:
     import pandas
 except (ImportError, AttributeError):
@@ -3108,8 +3112,15 @@ def test_list_rows_as_dataframe(client):
 
 
 @pytest.mark.skipif(pandas is None, reason="Requires `pandas`")
-@pytest.mark.skipif(pyarrow is None, reason="Requires `pyarrow`")
-def test_load_table_from_dataframe(client, to_delete):
+@pytest.mark.parametrize("parquet_engine", ["pyarrow", "fastparquet"])
+def test_load_table_from_dataframe(client, to_delete, parquet_engine):
+    if parquet_engine == "pyarrow" and pyarrow is None:
+        pytest.skip("Requires `pyarrow`")
+    if parquet_engine == "fastparquet" and fastparquet is None:
+        pytest.skip("Requires `fastparquet`")
+
+    pandas.set_option("io.parquet.engine", parquet_engine)
+
     dataset_id = "load_table_from_dataframe_{}".format(_millis())
     dataset = bigquery.Dataset(client.dataset(dataset_id))
     client.create_dataset(dataset)
diff --git a/bigquery/google/cloud/bigquery/client.py b/bigquery/google/cloud/bigquery/client.py
@@ -24,6 +24,7 @@
 import functools
 import gzip
 import os
+import tempfile
 import uuid
 
 import six
@@ -1124,10 +1125,10 @@ def load_table_from_dataframe(
         Raises:
             ImportError:
                 If a usable parquet engine cannot be found. This method
-                requires :mod:`pyarrow` to be installed.
+                requires :mod:`pyarrow` or :mod:`fastparquet` to be
+                installed.
         """
-        buffer = six.BytesIO()
-        dataframe.to_parquet(buffer)
+        job_id = _make_job_id(job_id, job_id_prefix)
 
         if job_config is None:
             job_config = job.LoadJobConfig()
@@ -1136,17 +1137,27 @@ def load_table_from_dataframe(
         if location is None:
             location = self.location
 
-        return self.load_table_from_file(
-            buffer,
-            destination,
-            num_retries=num_retries,
-            rewind=True,
-            job_id=job_id,
-            job_id_prefix=job_id_prefix,
-            location=location,
-            project=project,
-            job_config=job_config,
-        )
+        tmpfd, tmppath = tempfile.mkstemp(suffix="_job_{}.parquet".format(job_id[:8]))
+        os.close(tmpfd)
+
+        try:
+            dataframe.to_parquet(tmppath)
+
+            with open(tmppath, "rb") as parquet_file:
+                return self.load_table_from_file(
+                    parquet_file,
+                    destination,
+                    num_retries=num_retries,
+                    rewind=True,
+                    job_id=job_id,
+                    job_id_prefix=job_id_prefix,
+                    location=location,
+                    project=project,
+                    job_config=job_config,
+                )
+
+        finally:
+            os.remove(tmppath)
 
     def _do_resumable_upload(self, stream, metadata, num_retries):
         """Perform a resumable upload.
diff --git a/bigquery/noxfile.py b/bigquery/noxfile.py
@@ -124,7 +124,7 @@ def snippets(session):
         session.install('-e', local_dep)
     session.install('-e', os.path.join('..', 'storage'))
     session.install('-e', os.path.join('..', 'test_utils'))
-    session.install('-e', '.[pandas, pyarrow]')
+    session.install('-e', '.[pandas, pyarrow, fastparquet]')
 
     # Run py.test against the snippets tests.
     session.run(
diff --git a/bigquery/setup.py b/bigquery/setup.py
@@ -39,6 +39,7 @@
     # Exclude PyArrow dependency from Windows Python 2.7.
     'pyarrow: platform_system != "Windows" or python_version >= "3.4"':
         'pyarrow>=0.4.1',
+    'fastparquet': ['fastparquet', 'python-snappy'],
 }
 
 
diff --git a/bigquery/tests/unit/test_client.py b/bigquery/tests/unit/test_client.py
@@ -4658,17 +4658,15 @@ def test_load_table_from_dataframe(self):
             self.TABLE_REF,
             num_retries=_DEFAULT_NUM_RETRIES,
             rewind=True,
-            job_id=None,
+            job_id=mock.ANY,
             job_id_prefix=None,
             location=None,
             project=None,
             job_config=mock.ANY,
         )
 
         sent_file = load_table_from_file.mock_calls[0][1][1]
-        sent_bytes = sent_file.getvalue()
-        assert isinstance(sent_bytes, bytes)
-        assert len(sent_bytes) > 0
+        assert sent_file.closed
 
         sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
         assert sent_config.source_format == job.SourceFormat.PARQUET
@@ -4695,17 +4693,15 @@ def test_load_table_from_dataframe_w_client_location(self):
             self.TABLE_REF,
             num_retries=_DEFAULT_NUM_RETRIES,
             rewind=True,
-            job_id=None,
+            job_id=mock.ANY,
             job_id_prefix=None,
             location=self.LOCATION,
             project=None,
             job_config=mock.ANY,
         )
 
         sent_file = load_table_from_file.mock_calls[0][1][1]
-        sent_bytes = sent_file.getvalue()
-        assert isinstance(sent_bytes, bytes)
-        assert len(sent_bytes) > 0
+        assert sent_file.closed
 
         sent_config = load_table_from_file.mock_calls[0][2]["job_config"]
         assert sent_config.source_format == job.SourceFormat.PARQUET
@@ -4735,7 +4731,7 @@ def test_load_table_from_dataframe_w_custom_job_config(self):
             self.TABLE_REF,
             num_retries=_DEFAULT_NUM_RETRIES,
             rewind=True,
-            job_id=None,
+            job_id=mock.ANY,
             job_id_prefix=None,
             location=self.LOCATION,
             project=None,

Original file line number	Diff line number	Diff line change
`@@ -39,6 +39,7 @@`
`39`	`39`	`# Exclude PyArrow dependency from Windows Python 2.7.`
`40`	`40`	`'pyarrow: platform_system != "Windows" or python_version >= "3.4"':`
`41`	`41`	`'pyarrow>=0.4.1',`
	`42`	`+ 'fastparquet': ['fastparquet', 'python-snappy'],`
`42`	`43`	`}`
`43`	`44`
`44`	`45`