small tweaks

lincc-frameworks · Mar 18, 2024 · 949a85d · 949a85d
1 parent da39b86
commit 949a85d
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -18,7 +18,7 @@ requires-python = ">=3.9"
 dependencies = [
     'pandas',
     'numpy',
-    'dask>=2023.6.1,<2024.3.0', # We currently do not support dask/dask-expr
+    'dask',
     'dask[distributed]',
     'pyarrow',
     'pyvo',

diff --git a/src/tape/__init__.py b/src/tape/__init__.py
@@ -1,6 +1,7 @@
 from .analysis import *  # noqa
 from .ensemble import *  # noqa
-from .ensemble_frame import *  # noqa
+#from .ensemble_frame import *  # noqa
+from .expr import *  # noqa
 from .timeseries import *  # noqa
 from .ensemble_readers import *  # noqa
 from ._version import __version__  # noqa

diff --git a/tests/tape_tests/test_ensemble.py b/tests/tape_tests/test_ensemble.py
@@ -2158,7 +2158,7 @@ def test_select_random_timeseries(parquet_ensemble, repartition, seed):
     ens = parquet_ensemble
 
     if repartition:
-        ens.object = ens.object.repartition(3)
+        ens.object = ens.object.repartition(npartitions=3)
 
     ts = ens.select_random_timeseries(seed=seed)
 
@@ -2194,7 +2194,7 @@ def test_select_random_timeseries_empty_partitions(dask_client, all_empty):
     ens.from_source_dict(data_dict, column_mapper=colmap)
 
     # The single id will be in the last partition
-    ens.object = ens.object.repartition(5)
+    ens.object = ens.object.repartition(npartitions=5)
 
     # Remove the last partition, make sure we get the expected error when the
     # Object table has no IDs in any partition