Skip to content

Commit 924a960

Browse files
committed
[Python] Add tests for Date32 and Date64 array creation with masks
1 parent de6eb89 commit 924a960

File tree

1 file changed

+58
-2
lines changed

1 file changed

+58
-2
lines changed

python/pyarrow/tests/test_pandas.py

Lines changed: 58 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3278,14 +3278,70 @@ def test_error_sparse(self):
32783278

32793279

32803280
def test_safe_cast_from_float_with_nans_to_int():
3281-
# TODO(kszucs): write tests for creating Date32 and Date64 arrays, see
3282-
# ARROW-4258 and https://github.com/apache/arrow/pull/3395
32833281
values = pd.Series([1, 2, None, 4])
32843282
arr = pa.Array.from_pandas(values, type=pa.int32(), safe=True)
32853283
expected = pa.array([1, 2, None, 4], type=pa.int32())
32863284
assert arr.equals(expected)
32873285

32883286

3287+
def test_create_date32_and_date64_arrays_with_mask():
3288+
# Test Date32 array creation from Python list with mask
3289+
arr_date32 = pa.array([0, 0, 1, 2],
3290+
mask=[False, False, True, False],
3291+
type=pa.date32())
3292+
expected_date32 = pa.array([
3293+
date(1970, 1, 1),
3294+
date(1970, 1, 1),
3295+
None,
3296+
date(1970, 1, 3),
3297+
], type=pa.date32())
3298+
assert arr_date32.equals(expected_date32)
3299+
3300+
# Test Date32 array creation from Python dates
3301+
arr_date32_dates = pa.array([
3302+
date(2023, 1, 1),
3303+
date(2023, 1, 2),
3304+
None,
3305+
date(2023, 1, 4),
3306+
], type=pa.date32())
3307+
assert arr_date32_dates.null_count == 1
3308+
assert arr_date32_dates[2].as_py() is None
3309+
3310+
# Test Date64 array creation from Python list with mask
3311+
arr_date64 = pa.array([0, 86400000, 172800000, 259200000],
3312+
mask=[False, False, True, False],
3313+
type=pa.date64())
3314+
expected_date64 = pa.array([
3315+
date(1970, 1, 1),
3316+
date(1970, 1, 2),
3317+
None,
3318+
date(1970, 1, 4),
3319+
], type=pa.date64())
3320+
assert arr_date64.equals(expected_date64)
3321+
3322+
# Test Date64 array creation from Python dates
3323+
arr_date64_dates = pa.array([
3324+
date(2023, 1, 1),
3325+
date(2023, 1, 2),
3326+
None,
3327+
date(2023, 1, 4),
3328+
], type=pa.date64())
3329+
assert arr_date64_dates.null_count == 1
3330+
assert arr_date64_dates[2].as_py() is None
3331+
3332+
# Test Date32 with all nulls mask
3333+
arr_all_null = pa.array([0, 1, 2, 3],
3334+
mask=[True, True, True, True],
3335+
type=pa.date32())
3336+
assert arr_all_null.null_count == 4
3337+
3338+
# Test Date64 with no nulls
3339+
arr_no_null = pa.array([0, 86400000, 172800000],
3340+
mask=[False, False, False],
3341+
type=pa.date64())
3342+
assert arr_no_null.null_count == 0
3343+
3344+
32893345
def _fully_loaded_dataframe_example():
32903346
index = pd.MultiIndex.from_arrays([
32913347
pd.date_range('2000-01-01', periods=5).repeat(2),

0 commit comments

Comments
 (0)