Precompute string lengths sgkit-dev#12

sgkit_plink/pysnptools.py

@@ -1,11 +1,12 @@
 """PLINK 1.9 reader implementation"""
 from pathlib import Path
-from typing import Optional, Union
+from typing import Optional, Union, Mapping, Any
 
 import dask.array as da
 import dask.dataframe as dd
 import numpy as np
 from dask.dataframe import DataFrame
+from dask.array import Array
 from pysnptools.snpreader import Bed
 from xarray import Dataset
 
@@ -92,11 +93,26 @@ def close(self):
         self.bed._close_bed()  # pragma: no cover
 
 
-def _to_dict(df, dtype=None):
-    return {
-        c: df[c].to_dask_array(lengths=True).astype(dtype[c] if dtype else df[c].dtype)
-        for c in df
-    }
+def _max_str_len(arr: Array) -> Array:
+    return arr.map_blocks(
+        lambda s: np.char.str_len(s.astype(str)), dtype=np.int8
+    ).max()
+
+
+def _to_dict(df: DataFrame, dtype: Mapping[str, Any]=None):
+    arrs = {}
+    for c in df:
+        a = df[c].to_dask_array(lengths=True)
+        dt = df[c].dtype
+        if dtype:
+            dt = dtype[c]
+        kind = np.dtype(dt).kind
+        if kind in ['U', 'S']:
+            # Compute fixed-length string dtype for array
+            max_len = _max_str_len(a).compute()
+            dt = f"{kind}{max_len}"
+        arrs[c] = a.astype(dt)
+    return arrs
 
 
 def read_fam(path: PathType, sep: str = " ") -> DataFrame: