[hail] Promote localize_entries to public & tested

hail/python/hail/matrixtable.py

@@ -2533,6 +2533,63 @@ def _localize_entries(self, entries_field_name, cols_field_name) -> 'Table':
         return Table(CastMatrixToTable(
             self._mir, entries_field_name, cols_field_name))
 
+    @typecheck_method(entries_array_field_name=nullable(str),
+                      columns_array_field_name=nullable(str))
+    def localize_entries(self,
+                         entries_array_field_name=None,
+                         columns_array_field_name=None) -> 'Table':
+        """Represent this matrix as a table of entry-rows.
+
+        Examples
+        --------
+        Build a numpy ndarray from a small :class:`.MatrixTable`:
+
+        >>> mt = hl.utils.range_matrix_table(3,3)
+        >>> mt = mt.select_entries(x = mt.row_idx * mt.col_idx)
+        >>> mt.x.show()
+        +---------+---------+-------+
+        | row_idx | col_idx |     x |
+        +---------+---------+-------+
+        |   int32 |   int32 | int32 |
+        +---------+---------+-------+
+        |       0 |       2 |     0 |
+        |       0 |       0 |     0 |
+        |       0 |       1 |     0 |
+        |       1 |       0 |     0 |
+        |       1 |       1 |     1 |
+        |       1 |       2 |     2 |
+        |       2 |       2 |     4 |
+        |       2 |       1 |     2 |
+        |       2 |       0 |     0 |
+        +---------+---------+-------+
+        >>> t = mt.localize_entries('entry_structs', 'columns')
+        >>> t = t.select(entries = t.entry_structs.map(lambda entry: entry.x))
+        >>> np.array(t.entries.collect())
+        array([[0, 0, 0],
+               [0, 1, 2],
+               [0, 2, 4]])
+
+        Parameters
+        ----------
+        entries_array_field_name : :obj:`str`
+            The name of the table field containing the array of entry structs
+            for the given row
+        columns_array_field_name : :obj:`str`
+            The name of the global field containing the array of column structs
+
+        Returns
+        -------
+        :class:`.Table`
+        """
+        entries = entries_array_field_name or Env.get_uid()
+        cols = columns_array_field_name or Env.get_uid()
+        t = self._localize_entries(self, entries, cols)
+        if entries_array_field_name is None:
+            t = t.drop(entries)
+        if columns_array_field_name is None:
+            t = t.drop_globals(cols)
+        return t
+
     def _unfilter_entries(self):
         entry_ir = hl.cond(hl.is_defined(self.entry), self.entry, hl.struct(**self.entry))._ir
         return MatrixTable(MatrixMapEntries(self._mir, entry_ir))

hail/python/test/hail/matrixtable/test_matrix_table.py

@@ -5,6 +5,7 @@
 
 import hail as hl
 import hail.expr.aggregators as agg
+from hail.utils.java import Env
 from hail.utils.misc import new_temp_file
 from ..helpers import *
 
@@ -605,7 +606,7 @@ def test_read_stored_globals(self):
         self.assertTrue(ds.globals_table()._same(t))
 
     def test_codecs_matrix(self):
-        from hail.utils.java import Env, scala_object
+        from hail.utils.java import scala_object
         codecs = scala_object(Env.hail().io, 'CodecSpec').codecSpecs()
         ds = self.get_vds()
         temp = new_temp_file(suffix='hmt')
@@ -615,7 +616,7 @@ def test_codecs_matrix(self):
             self.assertTrue(ds._same(ds2))
 
     def test_codecs_table(self):
-        from hail.utils.java import Env, scala_object
+        from hail.utils.java import scala_object
         codecs = scala_object(Env.hail().io, 'CodecSpec').codecSpecs()
         rt = self.get_vds().rows()
         temp = new_temp_file(suffix='ht')
@@ -1046,3 +1047,35 @@ def test_agg_cols_group_by(self):
                  ]
         for aggregation, expected in tests:
             self.assertEqual(t.select_rows(result = aggregation).result.collect()[0], expected)
+
+    def localize_entries_with_both_none_is_rows_table(self):
+        mt = hl.utils.range_matrix_table(10, 10)
+        mt = mt.select_entries(x = mt.row_idx * mt.col_idx)
+        localized = mt.localize_entries(entries_array_field_name=None,
+                                        columns_array_field_name=None)
+        rows_table = mt.rows()
+        assert rows_table.collect() == localized.collect()
+        assert rows_table.globals_table().collect() == localized.globals_table().collect()
+
+    def localize_entries_with_none_cols_adds_no_globals(self):
+        mt = hl.utils.range_matrix_table(10, 10)
+        mt = mt.select_entries(x = mt.row_idx * mt.col_idx)
+        localized = mt.localize_entries(entries_array_field_name=Env.get_uid(),
+                                        columns_array_field_name=None)
+        assert mt.globals_table().collect() == localized.globals_table().collect()
+
+    def localize_entries_with_none_entries_changes_no_rows(self):
+        mt = hl.utils.range_matrix_table(10, 10)
+        mt = mt.select_entries(x = mt.row_idx * mt.col_idx)
+        localized = mt.localize_entries(entries_array_field_name=None,
+                                        columns_array_field_name=Env.get_uid())
+        rows_table = mt.rows()
+        assert rows_table.collect() == localized.collect()
+
+    def localize_entries_creates_arrays_of_entries_and_array_of_cols(self):
+        mt = hl.utils.range_matrix_table(10, 10)
+        mt = mt.select_entries(x = mt.row_idx * mt.col_idx)
+        localized = mt.localize_entries(entries_array_field_name='entries',
+                                        columns_array_field_name='cols')
+        assert [[x * y for x in range(0, 10)] for y in range(0, 10)] == localized.entries.collect()
+        assert range(0, 10) == localized.cols.collect()