[BEAM-11731][BEAM-10582] Allow pyarrow<4,numpy<1.21.0, improve pyarrow verification (#13892)

Author: TheNeuralBit

sdks/python/apache_beam/coders/row_coder_test.py

@@ -57,7 +57,7 @@
 class RowCoderTest(unittest.TestCase):
   JON_SNOW = Person(
       name="Jon Snow",
-      age=23,
+      age=np.int32(23),
       address=None,
       aliases=["crow", "wildling"],
       knows_javascript=False,
@@ -69,7 +69,7 @@ class RowCoderTest(unittest.TestCase):
       JON_SNOW,
       Person(
           "Daenerys Targaryen",
-          25,
+          np.int32(25),
           "Westeros",
           ["Mother of Dragons"],
           False,
@@ -79,7 +79,7 @@ class RowCoderTest(unittest.TestCase):
       ),
       Person(
           "Michael Bluth",
-          30,
+          np.int32(30),
           None, [],
           True,
           b"I've made a huge mistake", {},

sdks/python/apache_beam/dataframe/io_test.py

@@ -31,6 +31,7 @@
 
 import pandas as pd
 import pandas.testing
+import pytest
 from pandas.testing import assert_frame_equal
 from parameterized import parameterized
 
@@ -77,6 +78,11 @@ def test_read_write_csv(self):
     self.assertCountEqual(['a,b,c', '1,2,3', '3,4,7'],
                           set(self.read_all_lines(output + 'out.csv*')))
 
+  @pytest.mark.uses_pyarrow
+  def test_read_write_parquet(self):
+    self._run_read_write_test(
+        'parquet', {}, {}, dict(check_index=False), ['pyarrow'])
+
   @parameterized.expand([
       ('csv', dict(index_col=0)),
       ('csv', dict(index_col=0, splittable=True)),
@@ -100,7 +106,6 @@ def test_read_write_csv(self):
           dict(check_index=False)),
       ('html', dict(index_col=0), {}, {}, ['lxml']),
       ('excel', dict(index_col=0), {}, {}, ['openpyxl', 'xlrd']),
-      ('parquet', {}, {}, dict(check_index=False), ['pyarrow']),
   ])
   # pylint: disable=dangerous-default-value
   def test_read_write(
@@ -110,6 +115,18 @@ def test_read_write(
       write_kwargs={},
       check_options={},
       requires=()):
+    self._run_read_write_test(
+        format, read_kwargs, write_kwargs, check_options, requires)
+
+  # pylint: disable=dangerous-default-value
+  def _run_read_write_test(
+      self,
+      format,
+      read_kwargs={},
+      write_kwargs={},
+      check_options={},
+      requires=()):
+
     for module in requires:
       try:
         importlib.import_module(module)

sdks/python/apache_beam/dataframe/schemas.py

@@ -89,7 +89,7 @@
 
 # Generate type map (presented visually in the docstring)
 _BIDIRECTIONAL = [
-    (np.bool, np.bool),
+    (bool, bool),
     (np.int8, np.int8),
     (np.int16, np.int16),
     (np.int32, np.int32),
@@ -100,9 +100,9 @@
     (pd.Int64Dtype(), Optional[np.int64]),
     (np.float32, Optional[np.float32]),
     (np.float64, Optional[np.float64]),
-    (np.object, Any),
+    (object, Any),
     (pd.StringDtype(), Optional[str]),
-    (pd.BooleanDtype(), Optional[np.bool]),
+    (pd.BooleanDtype(), Optional[bool]),
 ]
 
 PANDAS_TO_BEAM = {
@@ -173,7 +173,7 @@ def generate_proxy(element_type):
     for name, typehint in fields:
       # Default to np.object. This is lossy, we won't be able to recover
       # the type at the output.
-      dtype = BEAM_TO_PANDAS.get(typehint, np.object)
+      dtype = BEAM_TO_PANDAS.get(typehint, object)
       proxy[name] = proxy[name].astype(dtype)
 
     return proxy

sdks/python/apache_beam/dataframe/schemas_test.py

@@ -75,8 +75,8 @@ def check_df_pcoll_equal(actual):
     ([375, 24, None, 10, 16], pd.Int64Dtype(), 'i64_nullable'),
     ([375., 24., None, 10., 16.], np.float64, 'f64'),
     ([375., 24., None, 10., 16.], np.float32, 'f32'),
-    ([True, False, True, True, False], np.bool, 'bool'),
-    (['Falcon', 'Ostrich', None, 3.14, 0], np.object, 'any'),
+    ([True, False, True, True, False], bool, 'bool'),
+    (['Falcon', 'Ostrich', None, 3.14, 0], object, 'any'),
     ([True, False, True, None, False], pd.BooleanDtype(), 'bool_nullable'),
     (['Falcon', 'Ostrich', None, 'Aardvark', 'Elephant'],
      pd.StringDtype(),

sdks/python/apache_beam/io/parquetio_test.py

@@ -28,6 +28,7 @@
 
 import hamcrest as hc
 import pandas
+import pytest
 from parameterized import param
 from parameterized import parameterized
 
@@ -62,6 +63,7 @@
 
 
 @unittest.skipIf(pa is None, "PyArrow is not installed.")
+@pytest.mark.uses_pyarrow
 class TestParquet(unittest.TestCase):
   @classmethod
   def setUpClass(cls):

sdks/python/apache_beam/testing/synthetic_pipeline_test.py

@@ -36,9 +36,11 @@
 from apache_beam.testing.util import equal_to
 
 try:
-  import numpy as np
+  import numpy  # pylint: disable=unused-import
 except ImportError:
-  np = None
+  NP_INSTALLED = False
+else:
+  NP_INSTALLED = True
 
 
 def input_spec(
@@ -60,7 +62,8 @@ def input_spec(
   }
 
 
-@unittest.skipIf(np is None, 'Synthetic source dependencies are not installed')
+@unittest.skipIf(
+    not NP_INSTALLED, 'Synthetic source dependencies are not installed')
 class SyntheticPipelineTest(unittest.TestCase):
 
   # pylint: disable=expression-not-assigned

sdks/python/pytest.ini

@@ -30,6 +30,8 @@ markers =
     # Tests using this marker conflict with the xdist plugin in some way, such
     # as enabling save_main_session.
     no_xdist: run without pytest-xdist plugin
+    # We run these tests with multiple major pyarrow versions (BEAM-11211)
+    uses_pyarrow: tests that utilize pyarrow in some way
 
 # Default timeout intended for unit tests.
 # If certain tests need a different value, please see the docs on how to

sdks/python/setup.py

@@ -143,12 +143,11 @@ def get_version():
     'hdfs>=2.1.0,<3.0.0',
     'httplib2>=0.8,<0.18.0',
     'mock>=1.0.1,<3.0.0',
-    # TODO(BEAM-11731): Support numpy 1.20.0
-    'numpy>=1.14.3,<1.20.0',
+    'numpy>=1.14.3,<1.21.0',
     'pymongo>=3.8.0,<4.0.0',
     'oauth2client>=2.0.1,<5',
     'protobuf>=3.12.2,<4',
-    'pyarrow>=0.15.1,<3.0.0',
+    'pyarrow>=0.15.1,<4.0.0',
     'pydot>=1.2.0,<2',
     'python-dateutil>=2.8.0,<3',
     'pytz>=2018.3',

sdks/python/test-suites/tox/common.gradle

@@ -34,7 +34,8 @@ project.task("preCommitPy${pythonVersionSuffix}") {
       // Generates coverage reports only once, in Py38, to remove duplicated work
       if (pythonVersionSuffix.equals('38')) {
           dependsOn = ["testPy38CloudCoverage", "testPy38Cython",
-              "testPy38pyarrow-0", "testPy38pyarrow-1", "testPy38pyarrow-2"]
+              "testPy38pyarrow-0", "testPy38pyarrow-1", "testPy38pyarrow-2",
+              "testPy38pyarrow-3"]
       } else {
           dependsOn = ["testPy${pythonVersionSuffix}Cloud", "testPy${pythonVersionSuffix}Cython"]
       }

sdks/python/test-suites/tox/py38/build.gradle

@@ -37,9 +37,11 @@ testPy38Cython.mustRunAfter testPython38, testPy38CloudCoverage
 toxTask "testPy38pyarrow-0", "py38-pyarrow-0"
 toxTask "testPy38pyarrow-1", "py38-pyarrow-1"
 toxTask "testPy38pyarrow-2", "py38-pyarrow-2"
+toxTask "testPy38pyarrow-3", "py38-pyarrow-3"
 test.dependsOn "testPy38pyarrow-0"
 test.dependsOn "testPy38pyarrow-1"
 test.dependsOn "testPy38pyarrow-2"
+test.dependsOn "testPy38pyarrow-3"
 
 toxTask "whitespacelint", "whitespacelint"