Skip to content

Commit

Permalink
feat(python): 30x speedup initialising Series from python range o…
Browse files Browse the repository at this point in the history
…bject (#5397)
  • Loading branch information
alexander-beedie committed Nov 2, 2022
1 parent 446050b commit ad678ca
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 4 deletions.
10 changes: 6 additions & 4 deletions py-polars/polars/internals/lazy_functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -1387,11 +1387,13 @@ def arange(
"""
low = pli.expr_to_lit_or_expr(low, str_to_lit=False)
high = pli.expr_to_lit_or_expr(high, str_to_lit=False)

if eager:
df = pli.DataFrame({"a": [1]})
return df.select(arange(low, high, step).alias("arange"))["arange"]

return (
pli.DataFrame()
.select(arange(low, high, step))
.to_series()
.rename("arange", in_place=True)
)
return pli.wrap_expr(pyarange(low._pyexpr, high._pyexpr, step))


Expand Down
9 changes: 9 additions & 0 deletions py-polars/polars/internals/series/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -206,8 +206,10 @@ def __init__(
)
elif isinstance(values, Series):
self._s = series_to_pyseries(name, values)

elif _PYARROW_TYPE(values) and isinstance(values, (pa.Array, pa.ChunkedArray)):
self._s = arrow_to_pyseries(name, values)

elif _NUMPY_TYPE(values) and isinstance(values, np.ndarray):
self._s = numpy_to_pyseries(name, values, strict, nan_to_null)
if values.dtype.type == np.datetime64:
Expand All @@ -223,6 +225,13 @@ def __init__(

if dtype is not None:
self._s = self.cast(dtype, strict=True)._s

elif isinstance(values, range):
self._s = (
pli.arange(values.start, values.stop, values.step, eager=True)
.rename(name, in_place=True)
._s
)
elif isinstance(values, Sequence):
self._s = sequence_to_pyseries(
name, values, dtype=dtype, strict=strict, dtype_if_empty=dtype_if_empty
Expand Down

0 comments on commit ad678ca

Please sign in to comment.