From c7cfb9fda64986132f68261f75d567020a9c3c67 Mon Sep 17 00:00:00 2001
From: Pranjalitem189 <pranjalitembhurnikar9@gmail.com>
Date: Sun, 5 Jul 2026 17:09:17 -0700
Subject: [PATCH] Improve missing tsfile dependency error

---
 src/datasets/packaged_modules/tsfile/tsfile.py       | 12 ++++++++++++
 .../test_tsfile_missing_dependency.py                | 11 +++++++++++
 2 files changed, 23 insertions(+)
 create mode 100644 tests/packaged_modules/test_tsfile_missing_dependency.py

diff --git a/src/datasets/packaged_modules/tsfile/tsfile.py b/src/datasets/packaged_modules/tsfile/tsfile.py
index 96712b4516e..0199ba30a57 100644
--- a/src/datasets/packaged_modules/tsfile/tsfile.py
+++ b/src/datasets/packaged_modules/tsfile/tsfile.py
@@ -53,9 +53,16 @@
 # Type helpers
 # ---------------------------------------------------------------------------
 
+def _require_tsfile():
+    try:
+        import tsfile  # noqa: F401
+    except ImportError as e:
+        raise ImportError("To support loading TSFile datasets, please install 'tsfile'.") from e
+
 
 def _arrow_type(ts_dtype, *, unit: str, tz: Optional[str]) -> pa.DataType:
     """Map a tsfile ``TSDataType`` to its Arrow representation."""
+    _require_tsfile()
     from tsfile.constants import TSDataType
 
     return {
@@ -86,6 +93,7 @@ def _promote_tsdatatype(a, b):
     if a == b:
         return a
 
+    _require_tsfile()
     from tsfile.constants import TSDataType
 
     table = {
@@ -301,6 +309,7 @@ def _generate_tables(self, files):
 
     def _scan_metadata(self, files) -> Optional[dict]:
         """Walk every file and unify table name, TAG columns, FIELD types."""
+        _require_tsfile()
         from tsfile.constants import TIME_COLUMN, ColumnCategory
 
         wanted_table = self._table
@@ -431,6 +440,7 @@ def _build_device_index(self, readers: dict):
         - ``file_meta``: maps each readable file to its per-file context
           (``tag_cols``, ``field_cols``, ``time_col``).
         """
+        _require_tsfile()
         from tsfile.constants import ColumnCategory
 
         device_to_files: dict[tuple, list[str]] = {}
@@ -541,6 +551,7 @@ def _read_device_from_file(self, reader, meta: dict, device_key: tuple) -> tuple
         field columns that this file owns *and* that the builder requested;
         callers fill missing fields with all-null contributions.
         """
+        _require_tsfile()
         from tsfile import tag_eq
 
         file_tag_cols: list[str] = meta["tag_cols"]
@@ -730,6 +741,7 @@ def _open_reader(file: str):
         segfaults. The 6-byte ``TsFile`` magic header is checked first to
         bail out cleanly.
         """
+        _require_tsfile()
         from tsfile import TsFileReader
 
         try:
diff --git a/tests/packaged_modules/test_tsfile_missing_dependency.py b/tests/packaged_modules/test_tsfile_missing_dependency.py
new file mode 100644
index 00000000000..d13560ee9b9
--- /dev/null
+++ b/tests/packaged_modules/test_tsfile_missing_dependency.py
@@ -0,0 +1,11 @@
+import pytest
+
+from datasets import load_dataset
+
+
+def test_tsfile_missing_dependency_message(tmp_path):
+    tsfile_path = tmp_path / "sample.tsfile"
+    tsfile_path.write_bytes(b"fake tsfile content")
+
+    with pytest.raises(ImportError, match=r"To support loading TSFile datasets, please install 'tsfile'\."):
+        load_dataset("tsfile", data_files=str(tsfile_path))
\ No newline at end of file