python lit allow numpy literals

pola-rs · Feb 1, 2022 · be32b14 · be32b14
1 parent 61146e8
commit be32b14
Show file tree

Hide file tree

Showing 4 changed files with 61 additions and 1 deletion.
diff --git a/py-polars/polars/convert.py b/py-polars/polars/convert.py
@@ -67,6 +67,8 @@ def from_records(
     """
     Construct a DataFrame from a numpy ndarray or sequence of sequences.
 
+    Note that this is slower than creating from columnar memory.
+
     Parameters
     ----------
     data : numpy ndarray or Sequence of sequences

diff --git a/py-polars/polars/internals/frame.py b/py-polars/polars/internals/frame.py
@@ -2827,6 +2827,51 @@ def upsample(
         Or combine them:
         "3d12h4m25s" # 3 days, 12 hours, 4 minutes, and 25 seconds
 
+        Examples
+        --------
+
+        Upsample a DataFrame by a certain interval.
+
+        >>> from datetime import datetime
+        >>> df = pl.DataFrame(
+        ...     {
+        ...         "time": [
+        ...             datetime(2021, 2, 1),
+        ...             datetime(2021, 4, 1),
+        ...             datetime(2021, 5, 1),
+        ...             datetime(2021, 6, 1),
+        ...         ],
+        ...         "groups": ["A", "B", "A", "B"],
+        ...         "values": [0, 1, 2, 3],
+        ...     }
+        ... )
+        >>> (
+        ...     df.upsample(
+        ...         time_column="time", every="1mo", by="groups", maintain_order=True
+        ...     ).select(pl.all().forward_fill())
+        ... )
+
+        shape: (7, 3)
+        ┌─────────────────────┬────────┬────────┐
+        │ time                ┆ groups ┆ values │
+        │ ---                 ┆ ---    ┆ ---    │
+        │ datetime[ns]        ┆ str    ┆ i64    │
+        ╞═════════════════════╪════════╪════════╡
+        │ 2021-02-01 00:00:00 ┆ A      ┆ 0      │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+        │ 2021-03-01 00:00:00 ┆ A      ┆ 0      │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+        │ 2021-04-01 00:00:00 ┆ A      ┆ 0      │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+        │ 2021-05-01 00:00:00 ┆ A      ┆ 2      │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+        │ 2021-04-01 00:00:00 ┆ B      ┆ 1      │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+        │ 2021-05-01 00:00:00 ┆ B      ┆ 1      │
+        ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+        │ 2021-06-01 00:00:00 ┆ B      ┆ 3      │
+        └─────────────────────┴────────┴────────┘
+
         """
         if by is None:
             by = []

diff --git a/py-polars/polars/internals/lazy_functions.py b/py-polars/polars/internals/lazy_functions.py
@@ -486,7 +486,9 @@ def tail(
 
 
 def lit(
-    value: Optional[Union[float, int, str, date, datetime, "pli.Series"]],
+    value: Optional[
+        Union[float, int, str, date, datetime, "pli.Series", np.ndarray, Any]
+    ],
     dtype: Optional[Type[DataType]] = None,
 ) -> "pli.Expr":
     """
@@ -561,6 +563,10 @@ def lit(
 
     if dtype:
         return pli.wrap_expr(pylit(value)).cast(dtype)
+    # numpy literals like np.float32(0)
+    # have an item
+    if hasattr(value, "item"):
+        value = value.item()  # type: ignore
     return pli.wrap_expr(pylit(value))
 
 

diff --git a/py-polars/tests/test_interop.py b/py-polars/tests/test_interop.py
@@ -172,3 +172,10 @@ def test_cat_to_pandas() -> None:
     df = df.with_columns(pl.all().cast(pl.Categorical))
     out = df.to_pandas()
     assert "category" in str(out["a"].dtype)
+
+
+def test_numpy_to_lit() -> None:
+    out = pl.select(pl.lit(np.array([1, 2, 3]))).to_series().to_list()
+    assert out == [1, 2, 3]
+    out = pl.select(pl.lit(np.float32(0))).to_series().to_list()
+    assert out == [0.0]