Use i8 for positions if needed

Author: benjeffery

bio2zarr/tskit.py

@@ -152,11 +152,17 @@ def generate_schema(
             f"{schema_instance.dimensions['samples'].chunk_size}"
         )
 
+        # Check if positions will fit in i4 (max ~2.1 billion)
+        max_position = 0
+        if self.ts.num_sites > 0:
+            max_position = np.max(self.ts.sites_position)
+        position_dtype = "i4" if max_position <= np.iinfo(np.int32).max else "i8"
+
         array_specs = [
             vcz.ZarrArraySpec(
                 source="position",
                 name="variant_position",
-                dtype="i4",
+                dtype=position_dtype,
                 dimensions=["variants"],
                 description="Position of each variant",
             ),

tests/test_ts.py

@@ -115,6 +115,40 @@ def no_individuals_ts(self, tmp_path):
         tree_sequence.dump(ts_path)
         return ts_path, tree_sequence
 
+    def test_position_dtype_selection(self, tmp_path):
+        tables = tskit.TableCollection(sequence_length=100)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        tables.sites.add_row(position=10, ancestral_state="A")
+        tables.sites.add_row(position=20, ancestral_state="C")
+        ts_small = tables.tree_sequence()
+        ts_path_small = tmp_path / "small_positions.trees"
+        ts_small.dump(ts_path_small)
+
+        tables = tskit.TableCollection(sequence_length=3_000_000_000)
+        tables.nodes.add_row(flags=tskit.NODE_IS_SAMPLE, time=0)
+        tables.sites.add_row(position=10, ancestral_state="A")
+        tables.sites.add_row(position=np.iinfo(np.int32).max + 1, ancestral_state="C")
+        ts_large = tables.tree_sequence()
+        ts_path_large = tmp_path / "large_positions.trees"
+        ts_large.dump(ts_path_large)
+
+        ind_nodes = np.array([[0], [1]])
+        format_obj_small = ts.TskitFormat(ts_path_small, ind_nodes)
+        schema_small = format_obj_small.generate_schema()
+
+        position_field = next(
+            f for f in schema_small.fields if f.name == "variant_position"
+        )
+        assert position_field.dtype == "i4"
+
+        format_obj_large = ts.TskitFormat(ts_path_large, ind_nodes)
+        schema_large = format_obj_large.generate_schema()
+
+        position_field = next(
+            f for f in schema_large.fields if f.name == "variant_position"
+        )
+        assert position_field.dtype == "i8"
+
     def test_initialization(self, simple_ts):
         ts_path, tree_sequence = simple_ts
         ind_nodes = np.array([[0, 1], [2, 3]])