Merge pull request #384 from Renumics/fix/buggy-pyarrow-take

Fix/buggy pyarrow table take
Renumics · Dec 5, 2023 · b3b8d48 · b3b8d48
2 parents cbc3826 + 5946efe
commit b3b8d48
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/renumics/spotlight_plugins/core/huggingface_datasource.py b/renumics/spotlight_plugins/core/huggingface_datasource.py
@@ -90,9 +90,9 @@ def get_column_values(
                 raw_values = self._dataset.data[column_name]
             else:
                 actual_indices = list(range(len(self._dataset)))[indices]
-                raw_values = self._dataset.data[column_name].take(actual_indices)
+                raw_values = self._dataset.data.fast_gather(actual_indices)[column_name]
         else:
-            raw_values = self._dataset.data[column_name].take(indices)
+            raw_values = self._dataset.data.fast_gather(indices)[column_name]
 
         feature = self._dataset.features[column_name]