Skip to content

Commit 2c1c5e3

Browse files
authored
fix: maintain tupleness of RecordForms when using _select_columns and _prune_columns (#3410)
* Don't add field names to copy if is_tuple * Added test
1 parent 270081a commit 2c1c5e3

File tree

2 files changed

+28
-2
lines changed

2 files changed

+28
-2
lines changed

src/awkward/forms/recordform.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,9 @@ def _prune_columns(self, is_inside_record_or_union: bool) -> Self | None:
127127
if not fields and is_inside_record_or_union:
128128
return None
129129
else:
130-
return self.copy(contents=contents, fields=fields)
130+
return self.copy(
131+
contents=contents, fields=None if self.is_tuple else fields
132+
)
131133

132134
def _select_columns(self, match_specifier: _SpecifierMatcher) -> Self:
133135
contents = []
@@ -145,7 +147,7 @@ def _select_columns(self, match_specifier: _SpecifierMatcher) -> Self:
145147
contents.append(next_content)
146148
fields.append(field)
147149

148-
return self.copy(contents=contents, fields=fields)
150+
return self.copy(contents=contents, fields=None if self.is_tuple else fields)
149151

150152
def _column_types(self):
151153
return sum((x._column_types() for x in self._contents), ())
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# BSD 3-Clause License; see https://github.com/scikit-hep/awkward/blob/main/LICENSE
2+
3+
from __future__ import annotations
4+
5+
import awkward as ak
6+
7+
8+
def test_select_columns():
9+
(records, records_tuple) = (
10+
ak.forms.RecordForm(
11+
[
12+
ak.forms.NumpyForm("int64"),
13+
ak.forms.NumpyForm("int64"),
14+
],
15+
None if is_tuple else ["x", "y"],
16+
)
17+
for is_tuple in (False, True)
18+
)
19+
20+
assert not records.is_tuple
21+
assert records_tuple.is_tuple
22+
23+
assert not records.select_columns("*").is_tuple
24+
assert records_tuple.select_columns("*").is_tuple

0 commit comments

Comments
 (0)