Skip to content

Commit a143259

Browse files
jrebackjreback
committed
add ohlc, last, nth
1 parent aa0a659 commit a143259

File tree

2 files changed

+422
-400
lines changed

2 files changed

+422
-400
lines changed

‎pandas/src/generate_code.py

+116-110
Original file line numberDiff line numberDiff line change
@@ -700,7 +700,7 @@ def group_last_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
700700
Only aggregates on axis=0
701701
'''
702702
cdef:
703-
Py_ssize_t i, j, N, K, lab
703+
Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
704704
%(dest_type2)s val, count
705705
ndarray[%(dest_type2)s, ndim=2] resx
706706
ndarray[int64_t, ndim=2] nobs
@@ -713,26 +713,27 @@ def group_last_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
713713
714714
N, K = (<object> values).shape
715715
716-
for i in range(N):
717-
lab = labels[i]
718-
if lab < 0:
719-
continue
716+
with nogil:
717+
for i in range(N):
718+
lab = labels[i]
719+
if lab < 0:
720+
continue
720721
721-
counts[lab] += 1
722-
for j in range(K):
723-
val = values[i, j]
722+
counts[lab] += 1
723+
for j in range(K):
724+
val = values[i, j]
724725
725-
# not nan
726-
if val == val:
727-
nobs[lab, j] += 1
728-
resx[lab, j] = val
726+
# not nan
727+
if val == val:
728+
nobs[lab, j] += 1
729+
resx[lab, j] = val
729730
730-
for i in range(len(counts)):
731-
for j in range(K):
732-
if nobs[i, j] == 0:
733-
out[i, j] = %(nan_val)s
734-
else:
735-
out[i, j] = resx[i, j]
731+
for i in range(ncounts):
732+
for j in range(K):
733+
if nobs[i, j] == 0:
734+
out[i, j] = %(nan_val)s
735+
else:
736+
out[i, j] = resx[i, j]
736737
"""
737738

738739
group_last_bin_template = """@cython.wraparound(False)
@@ -759,26 +760,27 @@ def group_last_bin_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
759760
760761
N, K = (<object> values).shape
761762
762-
b = 0
763-
for i in range(N):
764-
while b < ngroups - 1 and i >= bins[b]:
765-
b += 1
763+
with nogil:
764+
b = 0
765+
for i in range(N):
766+
while b < ngroups - 1 and i >= bins[b]:
767+
b += 1
766768
767-
counts[b] += 1
768-
for j in range(K):
769-
val = values[i, j]
769+
counts[b] += 1
770+
for j in range(K):
771+
val = values[i, j]
770772
771-
# not nan
772-
if val == val:
773-
nobs[b, j] += 1
774-
resx[b, j] = val
773+
# not nan
774+
if val == val:
775+
nobs[b, j] += 1
776+
resx[b, j] = val
775777
776-
for i in range(ngroups):
777-
for j in range(K):
778-
if nobs[i, j] == 0:
779-
out[i, j] = %(nan_val)s
780-
else:
781-
out[i, j] = resx[i, j]
778+
for i in range(ngroups):
779+
for j in range(K):
780+
if nobs[i, j] == 0:
781+
out[i, j] = %(nan_val)s
782+
else:
783+
out[i, j] = resx[i, j]
782784
"""
783785

784786
group_nth_bin_template = """@cython.wraparound(False)
@@ -805,27 +807,28 @@ def group_nth_bin_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
805807
806808
N, K = (<object> values).shape
807809
808-
b = 0
809-
for i in range(N):
810-
while b < ngroups - 1 and i >= bins[b]:
811-
b += 1
810+
with nogil:
811+
b = 0
812+
for i in range(N):
813+
while b < ngroups - 1 and i >= bins[b]:
814+
b += 1
812815
813-
counts[b] += 1
814-
for j in range(K):
815-
val = values[i, j]
816+
counts[b] += 1
817+
for j in range(K):
818+
val = values[i, j]
816819
817-
# not nan
818-
if val == val:
819-
nobs[b, j] += 1
820-
if nobs[b, j] == rank:
821-
resx[b, j] = val
820+
# not nan
821+
if val == val:
822+
nobs[b, j] += 1
823+
if nobs[b, j] == rank:
824+
resx[b, j] = val
822825
823-
for i in range(ngroups):
824-
for j in range(K):
825-
if nobs[i, j] == 0:
826-
out[i, j] = %(nan_val)s
827-
else:
828-
out[i, j] = resx[i, j]
826+
for i in range(ngroups):
827+
for j in range(K):
828+
if nobs[i, j] == 0:
829+
out[i, j] = %(nan_val)s
830+
else:
831+
out[i, j] = resx[i, j]
829832
"""
830833

831834
group_nth_template = """@cython.wraparound(False)
@@ -838,7 +841,7 @@ def group_nth_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
838841
Only aggregates on axis=0
839842
'''
840843
cdef:
841-
Py_ssize_t i, j, N, K, lab
844+
Py_ssize_t i, j, N, K, lab, ncounts = len(counts)
842845
%(dest_type2)s val, count
843846
ndarray[%(dest_type2)s, ndim=2] resx
844847
ndarray[int64_t, ndim=2] nobs
@@ -851,27 +854,28 @@ def group_nth_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
851854
852855
N, K = (<object> values).shape
853856
854-
for i in range(N):
855-
lab = labels[i]
856-
if lab < 0:
857-
continue
857+
with nogil:
858+
for i in range(N):
859+
lab = labels[i]
860+
if lab < 0:
861+
continue
858862
859-
counts[lab] += 1
860-
for j in range(K):
861-
val = values[i, j]
863+
counts[lab] += 1
864+
for j in range(K):
865+
val = values[i, j]
862866
863-
# not nan
864-
if val == val:
865-
nobs[lab, j] += 1
866-
if nobs[lab, j] == rank:
867-
resx[lab, j] = val
867+
# not nan
868+
if val == val:
869+
nobs[lab, j] += 1
870+
if nobs[lab, j] == rank:
871+
resx[lab, j] = val
868872
869-
for i in range(len(counts)):
870-
for j in range(K):
871-
if nobs[i, j] == 0:
872-
out[i, j] = %(nan_val)s
873-
else:
874-
out[i, j] = resx[i, j]
873+
for i in range(ncounts):
874+
for j in range(K):
875+
if nobs[i, j] == 0:
876+
out[i, j] = %(nan_val)s
877+
else:
878+
out[i, j] = resx[i, j]
875879
"""
876880

877881
group_add_template = """@cython.wraparound(False)
@@ -1720,48 +1724,50 @@ def group_ohlc_%(name)s(ndarray[%(dest_type2)s, ndim=2] out,
17201724
raise NotImplementedError("Argument 'values' must have only "
17211725
"one dimension")
17221726
else:
1723-
for i in range(N):
1724-
while b < ngroups - 1 and i >= bins[b]:
1725-
if not got_first:
1726-
out[b, 0] = NAN
1727-
out[b, 1] = NAN
1728-
out[b, 2] = NAN
1729-
out[b, 3] = NAN
1730-
else:
1731-
out[b, 0] = vopen
1732-
out[b, 1] = vhigh
1733-
out[b, 2] = vlow
1734-
out[b, 3] = vclose
1735-
b += 1
1736-
got_first = 0
17371727
1738-
counts[b] += 1
1739-
val = values[i, 0]
1728+
with nogil:
1729+
for i in range(N):
1730+
while b < ngroups - 1 and i >= bins[b]:
1731+
if not got_first:
1732+
out[b, 0] = NAN
1733+
out[b, 1] = NAN
1734+
out[b, 2] = NAN
1735+
out[b, 3] = NAN
1736+
else:
1737+
out[b, 0] = vopen
1738+
out[b, 1] = vhigh
1739+
out[b, 2] = vlow
1740+
out[b, 3] = vclose
1741+
b += 1
1742+
got_first = 0
17401743
1741-
# not nan
1742-
if val == val:
1743-
if not got_first:
1744-
got_first = 1
1745-
vopen = val
1746-
vlow = val
1747-
vhigh = val
1748-
else:
1749-
if val < vlow:
1744+
counts[b] += 1
1745+
val = values[i, 0]
1746+
1747+
# not nan
1748+
if val == val:
1749+
if not got_first:
1750+
got_first = 1
1751+
vopen = val
17501752
vlow = val
1751-
if val > vhigh:
17521753
vhigh = val
1753-
vclose = val
1754-
1755-
if not got_first:
1756-
out[b, 0] = NAN
1757-
out[b, 1] = NAN
1758-
out[b, 2] = NAN
1759-
out[b, 3] = NAN
1760-
else:
1761-
out[b, 0] = vopen
1762-
out[b, 1] = vhigh
1763-
out[b, 2] = vlow
1764-
out[b, 3] = vclose
1754+
else:
1755+
if val < vlow:
1756+
vlow = val
1757+
if val > vhigh:
1758+
vhigh = val
1759+
vclose = val
1760+
1761+
if not got_first:
1762+
out[b, 0] = NAN
1763+
out[b, 1] = NAN
1764+
out[b, 2] = NAN
1765+
out[b, 3] = NAN
1766+
else:
1767+
out[b, 0] = vopen
1768+
out[b, 1] = vhigh
1769+
out[b, 2] = vlow
1770+
out[b, 3] = vclose
17651771
"""
17661772

17671773
arrmap_template = """@cython.wraparound(False)

0 commit comments

Comments
 (0)