Skip to content

Commit

Permalink
apacheGH-38655: [C++] "iso_calendar" kernel returns incorrect results…
Browse files Browse the repository at this point in the history
… for array length > 32 (apache#39360)

### Rationale for this change

When defining `StructArray`'s field builders for `ISOCalendar` we don't pre-allocate memory and then use unsafe append. This causes the resulting array to be at most 32 rows long.

### What changes are included in this PR?

This introduces required memory pre-allocation in the `ISOCalendar` c++ kernel.

### Are these changes tested?

This adds a test for the Python wrapper.

### Are there any user-facing changes?

Fixes the behavior of `iso_calendar` kernel.
* Closes: apache#38655

Lead-authored-by: Rok Mihevc <rok@mihevc.org>
Co-authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
Signed-off-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
  • Loading branch information
rok and jorisvandenbossche committed Jan 23, 2024
1 parent eed53bb commit 7e9f265
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 1 deletion.
2 changes: 1 addition & 1 deletion cpp/src/arrow/compute/kernels/scalar_temporal_unary.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1510,7 +1510,7 @@ struct ISOCalendar {
for (int i = 0; i < 3; i++) {
field_builders.push_back(
checked_cast<BuilderType*>(struct_builder->field_builder(i)));
RETURN_NOT_OK(field_builders[i]->Reserve(1));
RETURN_NOT_OK(field_builders[i]->Reserve(in.length));
}
auto visit_null = [&]() { return struct_builder->AppendNull(); };
std::function<Status(typename InType::c_type arg)> visit_value;
Expand Down
13 changes: 13 additions & 0 deletions python/pyarrow/tests/test_compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -2263,6 +2263,19 @@ def test_extract_datetime_components():
_check_datetime_components(timestamps, timezone)


@pytest.mark.parametrize("unit", ["s", "ms", "us", "ns"])
def test_iso_calendar_longer_array(unit):
# https://github.com/apache/arrow/issues/38655
# ensure correct result for array length > 32
arr = pa.array([datetime.datetime(2022, 1, 2, 9)]*50, pa.timestamp(unit))
result = pc.iso_calendar(arr)
expected = pa.StructArray.from_arrays(
[[2021]*50, [52]*50, [7]*50],
names=['iso_year', 'iso_week', 'iso_day_of_week']
)
assert result.equals(expected)


@pytest.mark.pandas
@pytest.mark.skipif(sys.platform == "win32" and not util.windows_has_tzdata(),
reason="Timezone database is not installed on Windows")
Expand Down

0 comments on commit 7e9f265

Please sign in to comment.