Allow update of parquet categorical refs (#1510)

fsspec · Jan 23, 2024 · 0b06d92 · 0b06d92
1 parent ec13d27
commit 0b06d92
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/fsspec/implementations/http.py b/fsspec/implementations/http.py
@@ -123,7 +123,7 @@ def close_session(loop, session):
             try:
                 sync(loop, session.close, timeout=0.1)
                 return
-            except (TimeoutError, FSTimeoutError):
+            except (TimeoutError, FSTimeoutError, NotImplementedError):
                 pass
         connector = getattr(session, "_connector", None)
         if connector is not None:

diff --git a/fsspec/implementations/reference.py b/fsspec/implementations/reference.py
@@ -417,6 +417,11 @@ def write(self, field, record, base_url=None, storage_options=None):
             raws = np.full(self.record_size, np.nan, dtype="O")
         for j, data in partition.items():
             if isinstance(data, list):
+                if (
+                    str(paths.dtype) == "category"
+                    and data[0] not in paths.dtype.categories
+                ):
+                    paths = paths.add_categories(data[0])
                 paths[j] = data[0]
                 if len(data) > 1:
                     offsets[j] = data[1]