Support newer versions of PyArrow (#34892)

* Support newer versions of PyArrow

* add context for lower bound that would let us remove the pyarrow-hotfix

* smaller jump for pyarrow

* run postcommit

* tweak dynamic work rebalancing test

* add comment explaining unit test behavior

* yapf

Author: jrmccluskey

.github/trigger_files/beam_PostCommit_Python.json

@@ -1,6 +1,5 @@
 {
   "comment": "Modify this file in a trivial way to cause this test suite to run.",
-  "https://github.com/apache/beam/pull/32440": "test new datastream runner for batch"
-  "modification": 10
+  "modification": 11
 }
 

.github/trigger_files/beam_PostCommit_Python_Dependency.json

@@ -1,4 +1,4 @@
 {
     "comment": "Modify this file in a trivial way to cause this test suite to run",
-    "modification": 1
+    "modification": 0
   }

sdks/python/apache_beam/io/parquetio_test.py

@@ -487,7 +487,11 @@ def test_read_with_splitting_multiple_row_group(self):
     self._run_parquet_test(file_name, None, 10000, True, expected_result)
 
   def test_dynamic_work_rebalancing(self):
-    file_name = self._write_data(count=120, row_group_size=20)
+    # This test depends on count being sufficiently large + the ratio of
+    # count to row_group_size also being sufficiently large (but the required
+    # ratio to pass varies for values of row_group_size and, somehow, the
+    # version of pyarrow being tested against.)
+    file_name = self._write_data(count=280, row_group_size=20)
     source = _create_parquet_source(file_name)
 
     splits = [split for split in source.split(desired_bundle_size=float('inf'))]

sdks/python/setup.py

@@ -145,9 +145,10 @@ def cythonize(*args, **kwargs):
   pyarrow_dependency = ['']
 else:
   pyarrow_dependency = [
-      'pyarrow>=3.0.0,<17.0.0',
+      'pyarrow>=3.0.0,<19.0.0',
       # NOTE(https://github.com/apache/beam/issues/29392): We can remove this
       # once Beam increases the pyarrow lower bound to a version that fixes CVE.
+      # (lower bound >= 14.0.1)
       'pyarrow-hotfix<1'
   ]
 

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

@@ -85,6 +85,14 @@ toxTask "testPy39pyarrow-16", "py39-pyarrow-16", "${posargs}"
 test.dependsOn "testPy39pyarrow-16"
 postCommitPyDep.dependsOn "testPy39pyarrow-16"
 
+toxTask "testPy39pyarrow-17", "py39-pyarrow-17", "${posargs}"
+test.dependsOn "testPy39pyarrow-17"
+postCommitPyDep.dependsOn "testPy39pyarrow-17"
+
+toxTask "testPy39pyarrow-18", "py39-pyarrow-18", "${posargs}"
+test.dependsOn "testPy39pyarrow-18"
+postCommitPyDep.dependsOn "testPy39pyarrow-18"
+
 // Create a test task for each supported minor version of pandas
 toxTask "testPy39pandas-14", "py39-pandas-14", "${posargs}"
 test.dependsOn "testPy39pandas-14"

sdks/python/tox.ini

@@ -321,7 +321,7 @@ extras = test
 commands =
   bash {toxinidir}/scripts/pytest_validates_runner.sh {envname} {toxinidir}/apache_beam/runners/portability/prism_runner_test.py {posargs}
 
-[testenv:py{39,310}-pyarrow-{3,9,10,11,12,13,14,15,16}]
+[testenv:py{39,310}-pyarrow-{3,9,10,11,12,13,14,15,16,17,18}]
 deps =
   # As a courtesy to users, test against the oldest allowed version of Pyarrow.
   # We'd have to increase the pyarrow lower bound when Python 3.9 is deprecated.
@@ -339,6 +339,8 @@ deps =
   14: pyarrow>=14,<15
   15: pyarrow>=15,<16
   16: pyarrow>=16,<17
+  17: pyarrow>=17,<18
+  18: pyarrow>=18,<19
   numpy==1.26.4
 commands =
   # Log pyarrow and numpy version for debugging