Skip to content

Commit 6020b85

Browse files
authored
Merge 2b35b7a into 131a74f
2 parents 131a74f + 2b35b7a commit 6020b85

3 files changed

Lines changed: 178 additions & 30 deletions

File tree

‎astropy/io/votable/tests/test_tree.py‎

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,9 @@
1919
W54,
2020
W57,
2121
)
22-
from astropy.io.votable.table import parse
22+
from astropy.io.votable.table import parse, parse_single_table
2323
from astropy.io.votable.tree import MivotBlock, Resource, VOTableFile
24+
from astropy.table import Table
2425
from astropy.utils.data import get_pkg_data_filename
2526

2627

@@ -708,3 +709,86 @@ def test_votablefile_write(tmp_path):
708709

709710
# format kwarg should be accepted and ignored for API compatibility
710711
vot.write(str(out), format="votable")
712+
713+
714+
@pytest.mark.parametrize("tabledata_format", ["tabledata", "binary", "binary2"])
715+
def test_read_table_longer_than_first_allocation(tabledata_format):
716+
"""
717+
Test that parse on 1000-row table written in each serialization gives back
718+
every value and mask.
719+
720+
Reading starts with room for 512 rows, so this grows the result arrays and
721+
then trims them to the number of rows read.
722+
"""
723+
values = np.ma.array(np.arange(1000.0), mask=np.arange(1000) % 7 == 0)
724+
buffer = io.BytesIO()
725+
Table({"a": values}).write(
726+
buffer, format="votable", tabledata_format=tabledata_format
727+
)
728+
buffer.seek(0)
729+
730+
column = parse_single_table(buffer).array["a"]
731+
732+
assert np.array_equal(column.mask, values.mask)
733+
assert np.array_equal(column.filled(-1), values.filled(-1))
734+
735+
736+
@pytest.mark.parametrize("tabledata_format", ["tabledata", "binary", "binary2"])
737+
def test_read_variable_length_column_longer_than_first_allocation(tabledata_format):
738+
"""
739+
Test that parse on a 1000-row table with a variable-length int column gives
740+
back every row's values.
741+
"""
742+
votable = VOTableFile()
743+
resource = Resource()
744+
votable.resources.append(resource)
745+
table = tree.TableElement(votable)
746+
resource.tables.append(table)
747+
table.fields.append(tree.Field(votable, name="v", datatype="int", arraysize="*"))
748+
table.create_arrays(1000)
749+
# Fill the column with a different number of values in each row
750+
expected = [list(range(i % 5)) for i in range(1000)]
751+
for i, values in enumerate(expected):
752+
table.array.data["v"][i] = np.ma.array(
753+
values, mask=np.zeros(len(values), bool), dtype=np.int32
754+
)
755+
buffer = io.BytesIO()
756+
votable.to_xml(buffer, tabledata_format=tabledata_format)
757+
buffer.seek(0)
758+
759+
column = parse_single_table(buffer).array["v"]
760+
761+
assert [list(values) for values in column] == expected
762+
763+
764+
def _votable_with_nrows(tabledata_format, nrows):
765+
"""
766+
Return the bytes of a 1000-row VOTable in *tabledata_format* whose TABLE
767+
element falsely says it has *nrows* rows.
768+
"""
769+
values = np.ma.array(np.arange(1000.0), mask=np.arange(1000) % 7 == 0)
770+
buffer = io.BytesIO()
771+
Table({"a": values}).write(
772+
buffer, format="votable", tabledata_format=tabledata_format
773+
)
774+
content = buffer.getvalue().replace(b"<TABLE>", f'<TABLE nrows="{nrows}">'.encode())
775+
return content, values
776+
777+
778+
@pytest.mark.parametrize("nrows", [1, 1000, 5000])
779+
@pytest.mark.parametrize("tabledata_format", ["tabledata", "binary", "binary2"])
780+
def test_read_table_with_nrows(tabledata_format, nrows):
781+
"""
782+
Test that parse on a 1000-row table whose TABLE element has an nrows
783+
attribute of 1, 1000 or 5000 gives back exactly the 1000 values and masks.
784+
785+
Reading starts with room for nrows rows, so a too-small value grows the
786+
result arrays many times, the correct value never resizes them, and a
787+
too-large value only trims them at the end.
788+
"""
789+
content, values = _votable_with_nrows(tabledata_format, nrows)
790+
791+
column = parse_single_table(io.BytesIO(content)).array["a"]
792+
793+
assert np.array_equal(column.mask, values.mask)
794+
assert np.array_equal(column.filled(-1), values.filled(-1))

‎astropy/io/votable/tree.py‎

Lines changed: 89 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -126,17 +126,84 @@
126126
# FACTORY FUNCTIONS
127127

128128

129-
def _resize(masked, new_size):
129+
class _MaskedArrayBuilder:
130130
"""
131-
Masked arrays can not be resized inplace, and `np.resize` and
132-
`ma.resize` are both incompatible with structured arrays.
133-
Therefore, we do all this.
131+
Zeroed data and mask arrays for a table being read and grown in place.
132+
133+
Masked arrays can not be resized in place, but these plain arrays own
134+
their memory, so ``ndarray.resize`` can grow and trim them without
135+
holding a second copy of the table. Call `to_masked` once they are
136+
filled.
137+
138+
`resize` can move the arrays to new memory, so they never leave this
139+
class until `to_masked`, and nothing here keeps a view of them.
134140
"""
135-
new_array = ma.zeros((new_size,), dtype=masked.dtype)
136-
length = min(len(masked), new_size)
137-
new_array[:length] = masked[:length]
138141

139-
return new_array
142+
def __init__(self, dtype, nrows):
143+
self._values = np.zeros(nrows, dtype=dtype)
144+
self._mask = np.zeros(nrows, dtype=ma.make_mask_descr(dtype))
145+
146+
def __len__(self):
147+
return len(self._values)
148+
149+
def set_row(self, index, values, mask):
150+
"""
151+
Set one row.
152+
153+
Parameters
154+
----------
155+
index : int
156+
The row to set.
157+
values : tuple
158+
One value per field of the dtype, in field order, as returned
159+
by the column converters
160+
mask : tuple
161+
One mask flag per field, in the same order
162+
"""
163+
self._values[index] = values
164+
self._mask[index] = mask
165+
166+
def set_rows(self, start, values, masks):
167+
"""
168+
Set consecutive rows, starting at row *start*.
169+
170+
Parameters
171+
----------
172+
start : int
173+
The first row to set.
174+
values : list of tuple
175+
One tuple per row, each as the *values* of `set_row`.
176+
masks : list of tuple
177+
One tuple per row, each as the *mask* of `set_row`, and as
178+
many as there are in *values*.
179+
"""
180+
stop = start + len(values)
181+
self._values[start:stop] = values
182+
self._mask[start:stop] = masks
183+
184+
def resize(self, nrows):
185+
"""Grow or trim both arrays to *nrows* rows, keeping their contents."""
186+
# refcheck=False: from Python 3.14, numpy's reference count check
187+
# rejects a count of 2, which the attribute and this call already
188+
# make. It is safe because no views of these arrays exist.
189+
# Resizing while a view exists would leave the view pointing at
190+
# freed memory.
191+
self._values.resize(nrows, refcheck=False)
192+
self._mask.resize(nrows, refcheck=False)
193+
194+
def to_masked(self):
195+
"""
196+
Wrap the arrays in a masked array without copying them.
197+
198+
The builder can not be used afterwards, since resizing would free
199+
the memory the masked array uses.
200+
"""
201+
values, mask = self._values, self._mask
202+
del self._values, self._mask
203+
# keep_mask=False: for a structured dtype, MaskedArray starts from
204+
# its own all-False mask and, with keep_mask=True, ORs this one into
205+
# it, which allocates a second full-size mask. Unneeded here.
206+
return ma.MaskedArray(values, mask=mask, copy=False, keep_mask=False)
140207

141208

142209
def _lookup_by_attr_factory(attr, unique, iterator, element_name, doc):
@@ -2988,17 +3055,15 @@ def _parse_tabledata(self, iterator, colnumbers, config):
29883055
# allocation is by factors of 1.5.
29893056
invalid = config.get("invalid", "exception")
29903057

2991-
# Need to have only one reference so that we can resize the
2992-
# array
2993-
array = self.array
3058+
builder = _MaskedArrayBuilder(self.array.dtype, len(self.array))
29943059
del self.array
29953060

29963061
fields = self.all_fields
29973062
parsers = [field.converter.parse for field in fields]
29983063
binparsers = [field.converter.binparse for field in fields]
29993064

30003065
numrows = 0
3001-
alloc_rows = len(array)
3066+
alloc_rows = len(builder)
30023067
colnumbers_bits = [i in colnumbers for i in range(len(fields))]
30033068
row_default = [field.converter.default for field in self.fields]
30043069
mask_default = [True] * len(self.fields)
@@ -3076,10 +3141,9 @@ def _parse_tabledata(self, iterator, colnumbers, config):
30763141
if len(array_chunk) == chunk_size:
30773142
while numrows + chunk_size > alloc_rows:
30783143
alloc_rows = self._resize_strategy(alloc_rows)
3079-
if alloc_rows != len(array):
3080-
array = _resize(array, alloc_rows)
3081-
array[numrows : numrows + chunk_size] = array_chunk
3082-
array.mask[numrows : numrows + chunk_size] = mask_chunk
3144+
if alloc_rows != len(builder):
3145+
builder.resize(alloc_rows)
3146+
builder.set_rows(numrows, array_chunk, mask_chunk)
30833147
numrows += chunk_size
30843148
array_chunk = []
30853149
mask_chunk = []
@@ -3091,17 +3155,15 @@ def _parse_tabledata(self, iterator, colnumbers, config):
30913155
# put the last chunk values in there.
30923156
alloc_rows = numrows + len(array_chunk)
30933157

3094-
array = _resize(array, alloc_rows)
3095-
array[numrows:] = array_chunk
3096-
if alloc_rows != 0:
3097-
array.mask[numrows:] = mask_chunk
3158+
builder.resize(alloc_rows)
3159+
builder.set_rows(numrows, array_chunk, mask_chunk)
30983160
numrows += len(array_chunk)
30993161

31003162
if self.nrows is not None and self.nrows >= 0 and self.nrows != numrows:
31013163
warn_or_raise(W18, W18, (self.nrows, numrows), config, pos)
31023164
self._nrows = numrows
31033165

3104-
return array
3166+
return builder.to_masked()
31053167

31063168
def _get_binary_data_stream(self, iterator, config):
31073169
have_local_stream = False
@@ -3171,21 +3233,19 @@ def careful_read(length):
31713233
def _parse_binary(self, mode, iterator, colnumbers, config, pos):
31723234
careful_read = self._get_binary_data_stream(iterator, config)
31733235

3174-
# Need to have only one reference so that we can resize the
3175-
# array
3176-
array = self.array
3236+
builder = _MaskedArrayBuilder(self.array.dtype, len(self.array))
31773237
del self.array
31783238

31793239
fields = self.all_fields
31803240
binparsers = [field.converter.binparse for field in fields]
31813241

31823242
numrows = 0
3183-
alloc_rows = len(array)
3243+
alloc_rows = len(builder)
31843244
while True:
31853245
# Resize result arrays if necessary
31863246
if numrows >= alloc_rows:
31873247
alloc_rows = self._resize_strategy(alloc_rows)
3188-
array = _resize(array, alloc_rows)
3248+
builder.resize(alloc_rows)
31893249

31903250
row_data = []
31913251
row_mask_data = []
@@ -3230,11 +3290,11 @@ def _parse_binary(self, mode, iterator, colnumbers, config, pos):
32303290
row[j] = row_data[i]
32313291
row_mask[j] = row_mask_data[i]
32323292

3233-
array[numrows] = tuple(row)
3234-
array.mask[numrows] = tuple(row_mask)
3293+
builder.set_row(numrows, tuple(row), tuple(row_mask))
32353294
numrows += 1
32363295

3237-
return _resize(array, numrows)
3296+
builder.resize(numrows)
3297+
return builder.to_masked()
32383298

32393299
def _parse_fits(self, iterator, extnum, config):
32403300
for start, tag, data, pos in iterator:
Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
Reading a VOTable grows its result arrays in place, instead of copying the
2+
whole table each time they fill up and once more at the end. This lowers the
3+
peak memory needed to read a large table by up to the size of the table, and
4+
makes reading about 15% faster.

0 commit comments

Comments
 (0)