[Feature] Tabular data import/export (pandas, CSV, Parquet, JSON) (#1679)

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

pyproject.toml

@@ -55,6 +55,9 @@ tests = [
     "pytest-benchmark"
 ]
 h5 = ["h5py>=3.8"]
+pandas = ["pandas>=1.5"]
+parquet = ["pyarrow>=10.0"]
+tabular = ["pandas>=1.5", "pyarrow>=10.0"]
 dev = ["pybind11>=2.13", "ninja"]
 typecheck = ["mypy>=1.0.0"]
 onnx = ["onnx", "onnxscript", "onnxruntime"]

tensordict/__init__.py

@@ -28,9 +28,13 @@
     _default_is_leaf as default_is_leaf,
     _is_leaf_nontensor as is_leaf_nontensor,
     from_any,
+    from_csv,
     from_dict,
     from_h5,
+    from_json,
     from_namedtuple,
+    from_pandas,
+    from_parquet,
     from_struct_array,
     from_tuple,
     get_defaults_to_none,
@@ -113,10 +117,14 @@
     "LazyStackedTensorDictStore",
     "NestedKey",
     # Factory functions
+    "from_csv",
     "from_dict",
     "from_any",
     "from_h5",
+    "from_json",
     "from_namedtuple",
+    "from_pandas",
+    "from_parquet",
     "from_struct_array",
     "from_tuple",
     "from_dataclass",