Repository: astropy/astropy

Issue:
Table unicode sandwich follow-up: allow Table FITS reader to return bytestrings
Currently in Py3 one gets unicode character arrays when reading string data from a FITS file.  This translation is done in the the machinery within io.fits that [lazily auto-converts](https://github.com/astropy/astropy/blob/02c237987dd6f8bb1752ac5020ec47bff57dcb35/astropy/io/fits/fitsrec.py#L979) any string column to U-type for cols that can be decoded to ASCII.  See also [here](https://github.com/astropy/astropy/pull/5700#issuecomment-303688107).

After #5700 there is the opportunity to have `Table.read('file.fits')` return a Table with `'S'`-type columns, thus saving a factor of 4 in memory.

However, without implementing the unicode sandwich directly in `io.fits` (as @embray has prototyped), we have to be careful that direct use of HDU columns still returns unicode in Py3 while when used for reading a Table it would return bytes.  It isn't immediately obvious to me how to implement this dual-behavior (in terms of API), but I haven't looked too hard.

There is also the question of whether this should be the default or requires opt-ing in.

cc: @saimn 

Discussion:
astropy-bot[bot]: Hi there @astrofrog :wave: - thanks for the pull request! I'm just a friendly :robot: that checks for issues related to the changelog and making sure that this pull request is milestoned and labeled correctly. This is mainly intended for the maintainers, so if you are not a maintainer you can ignore this, and a maintainer will let you know if any action is required on your part :smiley:.

I see this is a work in progress pull request. I'll report back on the checks once the PR is ready for review.

*If there are any issues with this message, please report them [here](https://github.com/astropy/astropy-bot/issues).*
mhvk: @astrofrog - Is there a way you can check where the large peak in memory usage happens in the `data.view(self._data_type)` statement? In particular, after taking a view to a new class, that class's `__array_finalize__` is called, and `FITS_rec` has a non-trivial one (see https://github.com/astropy/astropy/blob/master/astropy/io/fits/fitsrec.py#L216, in particular starting L229), including some changes made in #4539.

Overall, I fear this may need a different look at the reference cycle that #4539 tried to address...
astrofrog: @mhvk - ah good point I can check ``__array_finalize__``
astrofrog: @mhvk - it turns out the memory increase is indeed in ``__array_finalize__``. More soon!
astrofrog: Basically it looks like the 2Gb memory appears in this line in ``__array_finalize__``:

```python
            self._coldefs = ColDefs(self)
```

I'm going to have to put this on pause for a few days or more. If anyone has any time to investigate this further, please feel free to!
astrofrog: I have some possible fixes coming...
taldcroft: @astrofrog - this looks amazing!

The scope of changes within FITS are about what I imagined from the last time I looked at this and looks reasonable.  The other changes to Time and Table seem sensible as well.

In real life you need to actually do something with the data.  Can you compare timing for these use cases:
- Access all of the data for one column
- Access all of the data for all columns (in a column-wise fashion)
- Access all of the data in a row-wise fashion
- Access first 100 rows in a row-wise fashion

By "access" I mean any cheap operation that demands getting the actual value.
taldcroft: @astrofrog - also I would suggest using a test data file that is at least 2.2 Gb in size.  My first attempt at chunked ASCII table reading used mmap (via the fast reader) and it turns out this blew up for files bigger than 2 Gb (which is the 32-bit limit).  I was doing this on a modern 64-bit mac and could not find anything 32-bit-ish in my setup, and this problem was the end of that implementation of chunked reading.

Also (sorry) this should be tested on Windows as well with a big file.
astrofrog: @taldcroft - thanks! I'll do all this :) Would you be able to take a look at the test failures to see if they are related to corner cases in the unicode sandwich comparison stuff?
mhvk: Main comment: I wonder if we cannot just use bytes everywhere, since `Column` sandwiches this anyway. Or is the main worry that people might use `view` on that and then get an unexpected result?

Beyond being not sure about the name and exact interpretation, this looks great!
astrofrog: @mhvk - you mean ``return_bytes`` should default to ``True``? If that's what you mean then the issue is that this would break things for people using astropy.io.fits directly.
mhvk: I guess what I'm trying to say is that the moment you translate a fits column to a table `Column`, there is, as far as I can see, no use in going for dtype `U`, since `Column` already sandwiches that. I also thought that was the only place where the conversion was done, but may well be wrong... Anyway, for conversion to `Column`, I think it makes sense to have no conversion to `U` be the default, but that can also be done by having `Table.read` pass on the relevant argument.
mhvk: @astrofrog - ah, yes, now I see that it was already what I thought made most sense!

For the name, since what we're doing is just skipping ascii decoding, maybe just call the argument `decode_ascii`? (with a default of `True`) Or `decode_bytes='ascii' (where `None` would be no decoding).

astrofrog: @taldcroft - I managed to isolate the unicode sandwich bug I mentioned higher up and have reported it in https://github.com/astropy/astropy/issues/6838
astrofrog: @taldcroft @mhvk - do you think you would be able to look into https://github.com/astropy/astropy/issues/6838 and https://github.com/astropy/astropy/issues/6839? (bugs in astropy.table) If not, I'll see if I can take a look. It would be nice to get this PR merged in for 3.0 :)
taldcroft: @astrofrog - OK will try to look at this soon.  I did investigate briefly, but got sidetracked.
astrofrog: I've rebased to include @taldcroft's fixes in https://github.com/astropy/astropy/pull/6899 (thanks!) to see if this now works. Once https://github.com/astropy/astropy/pull/6899 is merged, I'll rebase. 
astrofrog: I think my next fight is going to be trying to get FITS writing to not use such a ridiculous amount of memory... I'm trying to *create* a 2.5Gb file and running into issues.
astrofrog: I'll need this fix to pytest-openfiles to get this to pass on Travis: https://github.com/astropy/pytest-openfiles/pull/7
astrofrog: Ok so there are issues with remaining references to open files when memory mapping is enabled, and I can't yet figure out how to avoid that. This isn't going to be an issue for users, but just an issue for getting the test suite to pass. However, even without memory mapping, this PR is useful because it already avoids a lot of unnecessary copying and also avoids transforming the FITS string columns to unicode. I've turned off memory mapping in Table.read for now. Current results for reading a 100,000,000 row table with three columns - one with floats, one with ints, and one with 12-character strings (2.6Gb in total)

```
from astropy.table import Table
t = Table.read('many_rows.fits')
t.info('stats')
```

**Before**

![before](https://user-images.githubusercontent.com/314716/33295498-94022f64-d3cd-11e7-9dda-96141ccb9853.png)

**After**

![after](https://user-images.githubusercontent.com/314716/33295266-6adab800-d3cc-11e7-9ecd-f4198fa05c0c.png)

So even without memory mapping this is already a lot faster and uses less memory. The 'before' version uses 4x the size of the table file in memory, while the 'after' version uses 1.5x more at most. The 'after' version is also >10x faster.

I propose that we review this as-is then try and turn on memory mapping in a separate PR and try and resolve the residual open file issues there.

@taldcroft @mhvk - as mentioned above, there are still issues with passing bytes to Time without explicitly decoding.

mhvk: p.s. I like `character_as_bytes`
astrofrog: @mhvk - thanks for your review! I've addressed your comments

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 01e38237df..8d9420e22e 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -43,6 +43,19 @@ astropy.io.fits
 - Added ``ver`` attribute to set the ``EXTVER`` header keyword to ``ImageHDU``
   and ``TableHDU``. [#6454]
 
+- The performance for reading FITS tables has been significantly improved,
+  in particular for cases where the tables contain one or more string columns
+  and when done through ``Table.read``. [#6821]
+
+- The ``Table.read`` now supports a ``memmap=`` keyword argument to control
+  whether or not to use  memory mapping when reading the table. [#6821]
+
+- When reading FITS tables with ``fits.open``, a new keyword argument
+  ``character_as_bytes`` can be passed - when set to `True`, character columns
+  are returned as Numpy byte arrays (Numpy type S) while when set to `False`,
+  the same columns are decoded to Unicode strings (Numpy type U) which uses more
+  memory. [#6821]
+
 astropy.io.misc
 ^^^^^^^^^^^^^^^
 
@@ -237,6 +250,11 @@ astropy.table
   ``convert_bytestring_to_unicode`` and ``convert_unicode_to_bytestring``
   methods it was needed only for Python2 usage. [#6655]
 
+- When reading in FITS tables with ``Table.read``, string columns are now
+  represented using Numpy byte (dtype ``S``) arrays rather than Numpy
+  unicode arrays (dtype ``U``). The ``Column`` class then ensures the
+  bytes are automatically converted to string as needed. [#6821]
+
 astropy.tests
 ^^^^^^^^^^^^^
 
diff --git a/astropy/io/fits/column.py b/astropy/io/fits/column.py
index 5ea78c96fa..2a69cb8a32 100644
--- a/astropy/io/fits/column.py
+++ b/astropy/io/fits/column.py
@@ -1227,7 +1227,7 @@ class Column(NotifierMixin):
                         fsize = dims[-1]
                     else:
                         fsize = np.dtype(format.recformat).itemsize
-                    return chararray.array(array, itemsize=fsize)
+                    return chararray.array(array, itemsize=fsize, copy=False)
                 else:
                     return _convert_array(array, np.dtype(format.recformat))
             elif 'L' in format:
diff --git a/astropy/io/fits/connect.py b/astropy/io/fits/connect.py
index b5fbc03fa7..eef96b1b4f 100644
--- a/astropy/io/fits/connect.py
+++ b/astropy/io/fits/connect.py
@@ -8,7 +8,7 @@ from collections import OrderedDict
 
 from .. import registry as io_registry
 from ... import units as u
-from ...table import Table
+from ...table import Table, Column, MaskedColumn
 from ...utils.exceptions import AstropyUserWarning
 from . import HDUList, TableHDU, BinTableHDU, GroupsHDU
 from .column import KEYWORD_NAMES
@@ -63,7 +63,8 @@ def is_fits(origin, filepath, fileobj, *args, **kwargs):
         return False
 
 
-def read_table_fits(input, hdu=None, astropy_native=False):
+def read_table_fits(input, hdu=None, astropy_native=False, memmap=False,
+                    character_as_bytes=True):
     """
     Read a Table object from an FITS file
 
@@ -89,6 +90,20 @@ def read_table_fits(input, hdu=None, astropy_native=False):
     astropy_native : bool, optional
         Read in FITS columns as native astropy objects where possible instead
         of standard Table Column objects. Default is False.
+    memmap : bool, optional
+        Whether to use memory mapping, which accesses data on disk as needed. If
+        you are only accessing part of the data, this is often more efficient.
+        If you want to access all the values in the table, and you are able to
+        fit the table in memory, you may be better off leaving memory mapping
+        off. However, if your table would not fit in memory, you should set this
+        to `True`.
+    character_as_bytes : bool, optional
+        If `True`, string columns are stored as Numpy byte arrays (dtype ``S``)
+        and are converted on-the-fly to unicode strings when accessing
+        individual elements. If you need to use Numpy unicode arrays (dtype
+        ``U``) internally, you should set this to `False`, but note that this
+        will use more memory. If set to `False`, string columns will not be
+        memory-mapped even if ``memmap`` is `True`.
     """
 
     if isinstance(input, HDUList):
@@ -127,7 +142,8 @@ def read_table_fits(input, hdu=None, astropy_native=False):
 
     else:
 
-        hdulist = fits_open(input)
+        hdulist = fits_open(input, character_as_bytes=character_as_bytes,
+                            memmap=memmap)
 
         try:
             return read_table_fits(hdulist, hdu=hdu,
@@ -138,23 +154,33 @@ def read_table_fits(input, hdu=None, astropy_native=False):
     # Check if table is masked
     masked = any(col.null is not None for col in table.columns)
 
-    # Convert to an astropy.table.Table object
-    # Note: in future, it may make more sense to do this column-by-column,
+    # TODO: in future, it may make more sense to do this column-by-column,
     # rather than via the structured array.
-    t = Table(table.data, masked=masked)
 
-    # Copy over null values if needed
-    if masked:
-        for col in table.columns:
+    # In the loop below we access the data using data[col.name] rather than
+    # col.array to make sure that the data is scaled correctly if needed.
+    data = table.data
+
+    columns = []
+    for col in data.columns:
+
+        # Set column data
+        if masked:
+            column = MaskedColumn(data=data[col.name], name=col.name, copy=False)
             if col.null is not None:
-                t[col.name].set_fill_value(col.null)
-                t[col.name].mask[t[col.name] == col.null] = True
+                column.set_fill_value(col.null)
+                column.mask[column.data == col.null] = True
+        else:
+            column = Column(data=data[col.name], name=col.name, copy=False)
 
-    # Copy over units
-    for col in table.columns:
+        # Copy over units
         if col.unit is not None:
-            t[col.name].unit = u.Unit(
-                col.unit, format='fits', parse_strict='silent')
+            column.unit = u.Unit(col.unit, format='fits', parse_strict='silent')
+
+        columns.append(column)
+
+    # Create Table object
+    t = Table(columns, masked=masked, copy=False)
 
     # TODO: deal properly with unsigned integers
 
diff --git a/astropy/io/fits/fitsrec.py b/astropy/io/fits/fitsrec.py
index ba7487b6d2..59ce3323cb 100644
--- a/astropy/io/fits/fitsrec.py
+++ b/astropy/io/fits/fitsrec.py
@@ -152,6 +152,7 @@ class FITS_rec(np.recarray):
     """
 
     _record_type = FITS_record
+    _character_as_bytes = False
 
     def __new__(subtype, input):
         """
@@ -217,6 +218,9 @@ class FITS_rec(np.recarray):
         if obj is None:
             return
 
+        if isinstance(obj, FITS_rec):
+            self._character_as_bytes = obj._character_as_bytes
+
         if isinstance(obj, FITS_rec) and obj.dtype == self.dtype:
             self._converted = obj._converted
             self._heapoffset = obj._heapoffset
@@ -610,9 +614,9 @@ class FITS_rec(np.recarray):
     def __del__(self):
         try:
             del self._coldefs
-
             if self.dtype.fields is not None:
                 for col in self._col_weakrefs:
+
                     if col.array is not None:
                         col.array = col.array.copy()
 
@@ -964,8 +968,9 @@ class FITS_rec(np.recarray):
         elif _bool and field.dtype != bool:
             field = np.equal(field, ord('T'))
         elif _str:
-            with suppress(UnicodeDecodeError):
-                field = decode_ascii(field)
+            if not self._character_as_bytes:
+                with suppress(UnicodeDecodeError):
+                    field = decode_ascii(field)
 
         if dim:
             # Apply the new field item dimensions
@@ -1226,7 +1231,6 @@ class FITS_rec(np.recarray):
         # TODO: It would be nice if these string column formatting
         # details were left to a specialized class, as is the case
         # with FormatX and FormatP
-        print(format)
         if 'A' in format:
             _pc = '{:'
         else:
diff --git a/astropy/io/fits/hdu/hdulist.py b/astropy/io/fits/hdu/hdulist.py
index a7395b7b72..809a8d9095 100644
--- a/astropy/io/fits/hdu/hdulist.py
+++ b/astropy/io/fits/hdu/hdulist.py
@@ -114,6 +114,12 @@ def fitsopen(name, mode='readonly', memmap=None, save_backup=False,
             If `True`, image data is not scaled using BSCALE/BZERO values
             when read.
 
+        - **character_as_bytes** : bool
+
+            Whether to return bytes for string columns. By default this is `False`
+            and (unicode) strings are returned, but this does not respect memory
+            mapping and loads the whole column in memory when accessed.
+
         - **ignore_blank** : bool
 
             If `True`, the BLANK keyword is ignored if present.
diff --git a/astropy/io/fits/hdu/table.py b/astropy/io/fits/hdu/table.py
index c078b8f527..18e89d40e2 100644
--- a/astropy/io/fits/hdu/table.py
+++ b/astropy/io/fits/hdu/table.py
@@ -247,6 +247,10 @@ class _TableBaseHDU(ExtensionHDU, _TableLikeHDU):
         If not given or None, it defaults to the value of the ``EXTVER``
         card of the ``header`` or 1.
         (default: None)
+    character_as_bytes : bool
+        Whether to return bytes for string columns. By default this is `False`
+        and (unicode) strings are returned, but this does not respect memory
+        mapping and loads the whole column in memory when accessed.
     """
 
     _manages_own_heap = False
@@ -259,12 +263,17 @@ class _TableBaseHDU(ExtensionHDU, _TableLikeHDU):
     which perform their own heap maintenance.
     """
 
-    def __init__(self, data=None, header=None, name=None, uint=False, ver=None):
+    def __init__(self, data=None, header=None, name=None, uint=False, ver=None,
+                 character_as_bytes=False):
+
         super().__init__(data=data, header=header, name=name, ver=ver)
 
         if header is not None and not isinstance(header, Header):
             raise ValueError('header must be a Header object.')
+
         self._uint = uint
+        self._character_as_bytes = character_as_bytes
+
         if data is DELAYED:
             # this should never happen
             if header is None:
@@ -400,6 +409,7 @@ class _TableBaseHDU(ExtensionHDU, _TableLikeHDU):
     def data(self):
         data = self._get_tbdata()
         data._coldefs = self.columns
+        data._character_as_bytes = self._character_as_bytes
         # Columns should now just return a reference to the data._coldefs
         del self.columns
         return data
@@ -701,6 +711,10 @@ class TableHDU(_TableBaseHDU):
         If not given or None, it defaults to the value of the ``EXTVER``
         card of the ``header`` or 1.
         (default: None)
+    character_as_bytes : bool
+        Whether to return bytes for string columns. By default this is `False`
+        and (unicode) strings are returned, but this does not respect memory
+        mapping and loads the whole column in memory when accessed.
 
     """
 
@@ -713,8 +727,8 @@ class TableHDU(_TableBaseHDU):
     __format_RE = re.compile(
         r'(?P<code>[ADEFIJ])(?P<width>\d+)(?:\.(?P<prec>\d+))?')
 
-    def __init__(self, data=None, header=None, name=None, ver=None):
-        super().__init__(data, header, name=name, ver=ver)
+    def __init__(self, data=None, header=None, name=None, ver=None, character_as_bytes=False):
+        super().__init__(data, header, name=name, ver=ver, character_as_bytes=character_as_bytes)
 
     @classmethod
     def match_header(cls, header):
@@ -813,13 +827,18 @@ class BinTableHDU(_TableBaseHDU):
         If not given or None, it defaults to the value of the ``EXTVER``
         card of the ``header`` or 1.
         (default: None)
+    character_as_bytes : bool
+        Whether to return bytes for string columns. By default this is `False`
+        and (unicode) strings are returned, but this does not respect memory
+        mapping and loads the whole column in memory when accessed.
 
     """
 
     _extension = 'BINTABLE'
     _ext_comment = 'binary table extension'
 
-    def __init__(self, data=None, header=None, name=None, uint=False, ver=None):
+    def __init__(self, data=None, header=None, name=None, uint=False, ver=None,
+                 character_as_bytes=False):
         from ....table import Table
         if isinstance(data, Table):
             from ..convenience import table_to_hdu
@@ -829,7 +848,8 @@ class BinTableHDU(_TableBaseHDU):
             data = hdu.data
             header = hdu.header
 
-        super().__init__(data, header, name=name, uint=uint, ver=ver)
+        super().__init__(data, header, name=name, uint=uint, ver=ver,
+                         character_as_bytes=character_as_bytes)
 
     @classmethod
     def match_header(cls, header):
diff --git a/astropy/table/table.py b/astropy/table/table.py
index 34b59920cb..934e538494 100644
--- a/astropy/table/table.py
+++ b/astropy/table/table.py
@@ -1481,7 +1481,7 @@ class Table:
         except ValueError:
             raise ValueError("Column {0} does not exist".format(name))
 
-    def add_column(self, col, index=None, name=None, rename_duplicate=False):
+    def add_column(self, col, index=None, name=None, rename_duplicate=False, copy=True):
         """
         Add a new Column object ``col`` to the table.  If ``index``
         is supplied then insert column before ``index`` position
@@ -1498,6 +1498,8 @@ class Table:
             Column name
         rename_duplicate : bool
             Uniquify column name if it already exist. Default is False.
+        copy : bool
+            Make a copy of the new column. Default is True.
 
         Examples
         --------
@@ -1568,7 +1570,7 @@ class Table:
         if name is not None:
             name = (name,)
 
-        self.add_columns([col], [index], name, rename_duplicate=rename_duplicate)
+        self.add_columns([col], [index], name, copy=copy, rename_duplicate=rename_duplicate)
 
     def add_columns(self, cols, indexes=None, names=None, copy=True, rename_duplicate=False):
         """


Test patch:
diff --git a/astropy/io/fits/tests/test_connect.py b/astropy/io/fits/tests/test_connect.py
index 6ad0ac8c83..f1b65b64d0 100644
--- a/astropy/io/fits/tests/test_connect.py
+++ b/astropy/io/fits/tests/test_connect.py
@@ -1,4 +1,5 @@
 import os
+import gc
 import pathlib
 import warnings
 
@@ -147,6 +148,33 @@ class TestSingleTable:
         t = Table.read(hdu)
         assert equal_data(t, self.data)
 
+    def test_memmap(self, tmpdir):
+        filename = str(tmpdir.join('test_simple.fts'))
+        t1 = Table(self.data)
+        t1.write(filename, overwrite=True)
+        t2 = Table.read(filename, memmap=False)
+        t3 = Table.read(filename, memmap=True)
+        assert equal_data(t2, t3)
+        # To avoid issues with --open-files, we need to remove references to
+        # data that uses memory mapping and force the garbage collection
+        del t1, t2, t3
+        gc.collect()
+
+    @pytest.mark.parametrize('memmap', (False, True))
+    def test_character_as_bytes(self, tmpdir, memmap):
+        filename = str(tmpdir.join('test_simple.fts'))
+        t1 = Table(self.data)
+        t1.write(filename, overwrite=True)
+        t2 = Table.read(filename, character_as_bytes=False, memmap=memmap)
+        t3 = Table.read(filename, character_as_bytes=True, memmap=memmap)
+        assert t2['b'].dtype.kind == 'U'
+        assert t3['b'].dtype.kind == 'S'
+        assert equal_data(t2, t3)
+        # To avoid issues with --open-files, we need to remove references to
+        # data that uses memory mapping and force the garbage collection
+        del t1, t2, t3
+        gc.collect()
+
 
 class TestMultipleHDU:
 
@@ -250,7 +278,7 @@ def test_masking_regression_1795():
     assert np.all(t['c3'].mask == np.array([False, False]))
     assert np.all(t['c4'].mask == np.array([False, False]))
     assert np.all(t['c1'].data == np.array([1, 2]))
-    assert np.all(t['c2'].data == np.array(['abc', 'xy ']))
+    assert np.all(t['c2'].data == np.array([b'abc', b'xy ']))
     assert_allclose(t['c3'].data, np.array([3.70000007153, 6.6999997139]))
     assert np.all(t['c4'].data == np.array([False, True]))
 
diff --git a/astropy/table/tests/test_table.py b/astropy/table/tests/test_table.py
index 7359c2edce..967862b798 100644
--- a/astropy/table/tests/test_table.py
+++ b/astropy/table/tests/test_table.py
@@ -1234,7 +1234,7 @@ class TestConvertNumpyArray():
         for idx in range(len(arr.dtype)):
             assert arr.dtype[idx].byteorder != non_native_order
 
-        with fits.open(filename) as hdul:
+        with fits.open(filename, character_as_bytes=True) as hdul:
             data = hdul[1].data
             for colname in data.columns.names:
                 assert np.all(data[colname] == arr[colname])