Fix tabular imports and CI deps

Author: vmoens

.github/unittest/linux/scripts/environment.yml

@@ -22,4 +22,6 @@ dependencies:
     - ninja
     - numpy<2.0.0
     - mosaicml-streaming
+    - pandas
+    - pyarrow
     - redis

.github/unittest/linux/scripts/install.sh

@@ -53,7 +53,7 @@ printf "* Installing tensordict\n"
 # then install tensordict without resolving dependencies to avoid any solver changing
 # the PyTorch build (stable vs nightly).
 python -m pip install -U packaging pyvers importlib_metadata
-python -m pip install redis
+python -m pip install redis pandas pyarrow
 python -m pip install -e . --no-deps
 
 # smoke test

tensordict/base.py

@@ -57,6 +57,17 @@
 from tensordict._nestedkey import NestedKey
 from tensordict._tensorcollection import TensorCollection
 from tensordict.memmap import MemoryMappedTensor
+from tensordict.tabular import (
+    _columns_to_tensordict,
+    _dataframe_to_tensordict,
+    _read_csv,
+    _read_json,
+    _read_parquet,
+    _tensordict_to_dataframe,
+    _write_csv,
+    _write_json,
+    _write_parquet,
+)
 from tensordict.utils import (
     _as_context_manager,
     _CloudpickleWrapper,
@@ -15183,8 +15194,6 @@ def from_pandas(
                 device=None,
                 is_shared=False)
         """
-        from tensordict.tabular import _dataframe_to_tensordict
-
         if cls is TensorDictBase:
             from tensordict._td import TensorDict
 
@@ -15230,8 +15239,6 @@ def to_pandas(self, *, separator: str | None = None):
             1  1  0.0
             2  2  0.0
         """
-        from tensordict.tabular import _tensordict_to_dataframe
-
         return _tensordict_to_dataframe(self, separator=separator)
 
     @classmethod
@@ -15278,8 +15285,6 @@ def from_csv(
             >>> td = TensorDict.from_csv("data.csv")
             >>> td = TensorDict.from_csv("data.csv", separator=".", dtype=torch.float32)
         """
-        from tensordict.tabular import _columns_to_tensordict, _read_csv
-
         if cls is TensorDictBase:
             from tensordict._td import TensorDict
 
@@ -15315,8 +15320,6 @@ def to_csv(self, path, *, separator: str | None = None, **kwargs):
             **kwargs: Additional keyword arguments forwarded to
                 ``pandas.DataFrame.to_csv``.
         """
-        from tensordict.tabular import _write_csv
-
         _write_csv(self, path, separator=separator, **kwargs)
 
     @classmethod
@@ -15367,8 +15370,6 @@ def from_parquet(
             >>> td = TensorDict.from_parquet("data.parquet")
             >>> td = TensorDict.from_parquet("data.parquet", columns=["obs", "reward"])
         """
-        from tensordict.tabular import _columns_to_tensordict, _read_parquet
-
         if cls is TensorDictBase:
             from tensordict._td import TensorDict
 
@@ -15404,8 +15405,6 @@ def to_parquet(self, path, *, separator: str | None = None, **kwargs):
             **kwargs: Additional keyword arguments forwarded to the Parquet
                 writer.
         """
-        from tensordict.tabular import _write_parquet
-
         _write_parquet(self, path, separator=separator, **kwargs)
 
     @classmethod
@@ -15459,8 +15458,6 @@ def from_json(
             >>> td = TensorDict.from_json("data.json")
             >>> td = TensorDict.from_json("data.jsonl", lines=True)
         """
-        from tensordict.tabular import _columns_to_tensordict, _read_json
-
         if cls is TensorDictBase:
             from tensordict._td import TensorDict
 
@@ -15503,8 +15500,6 @@ def to_json(
             **kwargs: Additional keyword arguments forwarded to the JSON
                 writer.
         """
-        from tensordict.tabular import _write_json
-
         _write_json(self, path, separator=separator, lines=lines, **kwargs)
 
     def to_h5(

tensordict/tabular.py

@@ -8,22 +8,21 @@
 
 from __future__ import annotations
 
+import importlib.util
 from pathlib import Path
 from typing import Any
 
 import numpy as np
 import torch
+from tensordict._tensorcollection import TensorCollection
+from tensordict.utils import is_non_tensor
 
 
 def _has_pandas() -> bool:
-    import importlib.util
-
     return importlib.util.find_spec("pandas") is not None
 
 
 def _has_pyarrow() -> bool:
-    import importlib.util
-
     return importlib.util.find_spec("pyarrow") is not None
 
 
@@ -51,11 +50,9 @@ def _unflatten_columns(flat_dict: dict, separator: str) -> dict:
 
 def _flatten_keys(td, separator: str) -> dict[str, Any]:
     """Flatten a TensorDict into a dict with separated key names."""
-    from tensordict.base import _is_tensor_collection, is_non_tensor
-
     result = {}
     for key, value in td.items():
-        if _is_tensor_collection(type(value)) and not is_non_tensor(value):
+        if isinstance(value, TensorCollection) and not is_non_tensor(value):
             sub = _flatten_keys(value, separator)
             for sub_key, sub_val in sub.items():
                 result[f"{key}{separator}{sub_key}"] = sub_val
@@ -132,14 +129,12 @@ def _tensordict_to_dataframe(td, *, separator: str | None):
     """Convert a TensorDict to a pandas DataFrame."""
     import pandas as pd
 
-    from tensordict.base import _is_tensor_collection, is_non_tensor
-
     if separator is not None:
         flat = _flatten_keys(td, separator)
     else:
         flat = {}
         for key, value in td.items():
-            if _is_tensor_collection(type(value)) and not is_non_tensor(value):
+            if isinstance(value, TensorCollection) and not is_non_tensor(value):
                 raise ValueError(
                     f"Nested TensorDict at key '{key}' requires a separator parameter "
                     "to flatten to DataFrame columns. Use to_pandas(separator='.')."
@@ -319,8 +314,6 @@ def _write_json(td, path, separator: str | None, lines: bool = False, **kwargs):
     else:
         import json
 
-        from tensordict.base import is_non_tensor
-
         if separator is not None:
             flat = _flatten_keys(td, separator)
         else:

test/test_tabular.py

@@ -9,7 +9,16 @@
 import pytest
 import torch
 
-from tensordict import tensorclass, TensorDict, TensorDictBase
+from tensordict import (
+    from_csv,
+    from_json,
+    from_pandas,
+    from_parquet,
+    tensorclass,
+    TensorDict,
+    TensorDictBase,
+)
+from tensordict.tensorclass import NonTensorData
 
 _has_pandas = importlib.util.find_spec("pandas") is not None
 _has_pyarrow = importlib.util.find_spec("pyarrow") is not None
@@ -42,7 +51,6 @@ def test_string_columns(self):
         assert td.batch_size == torch.Size([3])
         assert td["age"].dtype == torch.int64
         name_val = td["name"]
-        from tensordict.tensorclass import NonTensorData
 
         assert isinstance(name_val, NonTensorData) or hasattr(name_val, "tolist")
 
@@ -133,8 +141,6 @@ def test_categorical_columns(self):
     def test_module_level_function(self):
         import pandas as pd
 
-        from tensordict import from_pandas
-
         df = pd.DataFrame({"x": [1, 2, 3]})
         td = from_pandas(df)
         assert td.batch_size == torch.Size([3])
@@ -245,8 +251,6 @@ def test_csv_roundtrip(self, tmp_path):
     def test_module_level_function(self, tmp_path):
         import pandas as pd
 
-        from tensordict import from_csv
-
         csv_path = tmp_path / "test.csv"
         pd.DataFrame({"x": [1, 2]}).to_csv(csv_path, index=False)
         td = from_csv(csv_path)
@@ -307,8 +311,6 @@ def test_module_level_function(self, tmp_path):
         import pyarrow as pa
         import pyarrow.parquet as pq
 
-        from tensordict import from_parquet
-
         path = tmp_path / "test.parquet"
         table = pa.table({"x": [1, 2]})
         pq.write_table(table, str(path))
@@ -349,8 +351,6 @@ def test_to_json_lines(self, tmp_path):
         assert len(lines) == 3
 
     def test_module_level_function(self, tmp_path):
-        from tensordict import from_json
-
         path = tmp_path / "test.json"
         path.write_text(json.dumps([{"a": 1}, {"a": 2}]))
         td = from_json(path)
@@ -410,7 +410,7 @@ def test_csv_roundtrip(self, tmp_path):
         tc2 = TabularTensorClass.from_csv(csv_path)
         assert isinstance(tc2, TabularTensorClass)
         assert (tc2.x == tc.x).all()
-        assert torch.allclose(tc2.y, tc.y)
+        assert torch.allclose(tc2.y.to(tc.y.dtype), tc.y)
 
     def test_json_roundtrip(self, tmp_path):
         path = tmp_path / "tensorclass.json"
@@ -423,7 +423,7 @@ def test_json_roundtrip(self, tmp_path):
         tc2 = TabularTensorClass.from_json(path)
         assert isinstance(tc2, TabularTensorClass)
         assert (tc2.x == tc.x).all()
-        assert torch.allclose(tc2.y, tc.y)
+        assert torch.allclose(tc2.y.to(tc.y.dtype), tc.y)
 
     @pytest.mark.skipif(not _has_pyarrow, reason="pyarrow not found")
     def test_parquet_roundtrip(self, tmp_path):
@@ -437,4 +437,4 @@ def test_parquet_roundtrip(self, tmp_path):
         tc2 = TabularTensorClass.from_parquet(path)
         assert isinstance(tc2, TabularTensorClass)
         assert (tc2.x == tc.x).all()
-        assert torch.allclose(tc2.y, tc.y)
+        assert torch.allclose(tc2.y.to(tc.y.dtype), tc.y)