quarter expression (#3797)

polars/polars-lazy/src/dsl/dt.rs

@@ -69,16 +69,26 @@ impl DateLikeNameSpace {
         let function = move |s: Series| s.year().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("year")
+            .with_fmt("dt.year")
     }
 
     /// Get the month of a Date/Datetime
     pub fn month(self) -> Expr {
         let function = move |s: Series| s.month().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("month")
+            .with_fmt("dt.month")
     }
+
+    /// Extract quarter from underlying NaiveDateTime representation.
+    /// Quarters range from 1 to 4.
+    pub fn quarter(self) -> Expr {
+        let function = move |s: Series| s.quarter().map(|ca| ca.into_series());
+        self.0
+            .map(function, GetOutput::from_type(DataType::UInt32))
+            .with_fmt("dt.quarter")
+    }
+
     /// Extract the week from the underlying Date representation.
     /// Can be performed on Date and Datetime
 
@@ -88,7 +98,7 @@ impl DateLikeNameSpace {
         let function = move |s: Series| s.week().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("week")
+            .with_fmt("dt.week")
     }
 
     /// Extract the week day from the underlying Date representation.
@@ -99,7 +109,7 @@ impl DateLikeNameSpace {
         let function = move |s: Series| s.weekday().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("weekday")
+            .with_fmt("dt.weekday")
     }
 
     /// Get the month of a Date/Datetime
@@ -114,36 +124,36 @@ impl DateLikeNameSpace {
         let function = move |s: Series| s.ordinal_day().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("ordinal_day")
+            .with_fmt("dt.ordinal_day")
     }
     /// Get the hour of a Datetime/Time64
     pub fn hour(self) -> Expr {
         let function = move |s: Series| s.hour().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("hour")
+            .with_fmt("dt.hour")
     }
     /// Get the minute of a Datetime/Time64
     pub fn minute(self) -> Expr {
         let function = move |s: Series| s.minute().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("minute")
+            .with_fmt("dt.minute")
     }
 
     /// Get the second of a Datetime/Time64
     pub fn second(self) -> Expr {
         let function = move |s: Series| s.second().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("second")
+            .with_fmt("dt.second")
     }
     /// Get the nanosecond of a Time64
     pub fn nanosecond(self) -> Expr {
         let function = move |s: Series| s.nanosecond().map(|ca| ca.into_series());
         self.0
             .map(function, GetOutput::from_type(DataType::UInt32))
-            .with_fmt("nanosecond")
+            .with_fmt("dt.nanosecond")
     }
 
     pub fn timestamp(self, tu: TimeUnit) -> Expr {
@@ -152,6 +162,6 @@ impl DateLikeNameSpace {
                 move |s| s.timestamp(tu).map(|ca| ca.into_series()),
                 GetOutput::from_type(DataType::Int64),
             )
-            .with_fmt("timestamp")
+            .with_fmt("dt.timestamp")
     }
 }

polars/polars-time/src/chunkedarray/date.rs

@@ -16,6 +16,13 @@ pub trait DateMethods {
     /// Returns the year number in the calendar date.
     fn year(&self) -> Int32Chunked;
 
+    /// Extract month from underlying NaiveDateTime representation.
+    /// Quarters range from 1 to 4.
+    fn quarter(&self) -> UInt32Chunked {
+        let months = self.month();
+        months_to_quarters(months)
+    }
+
     /// Extract month from underlying NaiveDateTime representation.
     /// Returns the month number starting from 1.
     ///

polars/polars-time/src/chunkedarray/datetime.rs

@@ -40,6 +40,13 @@ pub trait DatetimeMethods: AsDatetime {
         cast_and_apply(self.as_datetime(), temporal::year)
     }
 
+    /// Extract quarter from underlying NaiveDateTime representation.
+    /// Quarters range from 1 to 4.
+    fn quarter(&self) -> UInt32Chunked {
+        let months = self.month();
+        months_to_quarters(months)
+    }
+
     /// Extract month from underlying NaiveDateTime representation.
     /// Returns the month number starting from 1.
     ///

polars/polars-time/src/chunkedarray/mod.rs

@@ -29,3 +29,9 @@ pub use utf8::Utf8Methods;
 pub fn unix_time() -> NaiveDateTime {
     NaiveDateTime::from_timestamp(0, 0)
 }
+
+// a separate function so that it is not compiled twice
+pub(crate) fn months_to_quarters(mut ca: UInt32Chunked) -> UInt32Chunked {
+    ca.apply_mut(|month| (month + 2) / 3);
+    ca
+}

polars/polars-time/src/series/mod.rs

@@ -243,31 +243,46 @@ pub trait TemporalMethods: AsSeries {
     }
 
     /// Extract month from underlying NaiveDateTime representation.
-    /// Returns the month number starting from 1.
-    ///
-    /// The return value ranges from 1 to 12.
-    fn month(&self) -> Result<UInt32Chunked> {
+    /// Returns the year number in the calendar date.
+    fn year(&self) -> Result<Int32Chunked> {
         let s = self.as_series();
         match s.dtype() {
             #[cfg(feature = "dtype-date")]
-            DataType::Date => s.date().map(|ca| ca.month()),
+            DataType::Date => s.date().map(|ca| ca.year()),
             #[cfg(feature = "dtype-datetime")]
-            DataType::Datetime(_, _) => s.datetime().map(|ca| ca.month()),
+            DataType::Datetime(_, _) => s.datetime().map(|ca| ca.year()),
+            _ => Err(PolarsError::InvalidOperation(
+                format!("operation not supported on dtype {:?}", s.dtype()).into(),
+            )),
+        }
+    }
+
+    /// Extract quarter from underlying NaiveDateTime representation.
+    /// Quarters range from 1 to 4.
+    fn quarter(&self) -> Result<UInt32Chunked> {
+        let s = self.as_series();
+        match s.dtype() {
+            #[cfg(feature = "dtype-date")]
+            DataType::Date => s.date().map(|ca| ca.quarter()),
+            #[cfg(feature = "dtype-datetime")]
+            DataType::Datetime(_, _) => s.datetime().map(|ca| ca.quarter()),
             _ => Err(PolarsError::InvalidOperation(
                 format!("operation not supported on dtype {:?}", s.dtype()).into(),
             )),
         }
     }
 
     /// Extract month from underlying NaiveDateTime representation.
-    /// Returns the year number in the calendar date.
-    fn year(&self) -> Result<Int32Chunked> {
+    /// Returns the month number starting from 1.
+    ///
+    /// The return value ranges from 1 to 12.
+    fn month(&self) -> Result<UInt32Chunked> {
         let s = self.as_series();
         match s.dtype() {
             #[cfg(feature = "dtype-date")]
-            DataType::Date => s.date().map(|ca| ca.year()),
+            DataType::Date => s.date().map(|ca| ca.month()),
             #[cfg(feature = "dtype-datetime")]
-            DataType::Datetime(_, _) => s.datetime().map(|ca| ca.year()),
+            DataType::Datetime(_, _) => s.datetime().map(|ca| ca.month()),
             _ => Err(PolarsError::InvalidOperation(
                 format!("operation not supported on dtype {:?}", s.dtype()).into(),
             )),

py-polars/docs/source/reference/expression.rst

@@ -283,6 +283,7 @@ The following methods are available under the `expr.dt` attribute.
     ExprDateTimeNameSpace.nanosecond
     ExprDateTimeNameSpace.nanoseconds
     ExprDateTimeNameSpace.ordinal_day
+    ExprDateTimeNameSpace.quarter
     ExprDateTimeNameSpace.second
     ExprDateTimeNameSpace.seconds
     ExprDateTimeNameSpace.strftime

py-polars/docs/source/reference/series.rst

@@ -233,6 +233,7 @@ The following methods are available under the `Series.dt` attribute.
     DateTimeNameSpace.nanosecond
     DateTimeNameSpace.nanoseconds
     DateTimeNameSpace.ordinal_day
+    DateTimeNameSpace.quarter
     DateTimeNameSpace.second
     DateTimeNameSpace.seconds
     DateTimeNameSpace.strftime

py-polars/polars/internals/expr.py

@@ -5033,6 +5033,19 @@ def year(self) -> Expr:
         """
         return wrap_expr(self._pyexpr.year())
 
+    def quarter(self) -> Expr:
+        """
+        Extract quarter from underlying Date representation.
+        Can be performed on Date and Datetime.
+
+        Returns the quarter ranging from 1 to 4.
+
+        Returns
+        -------
+        Quarter as UInt32
+        """
+        return wrap_expr(self._pyexpr.quarter())
+
     def month(self) -> Expr:
         """
         Extract month from underlying Date representation.

py-polars/polars/internals/frame.py

@@ -1761,6 +1761,10 @@ def __getitem__(
             if isinstance(item[0], str):
                 return self._from_pydf(self._df.select(item))
             if item.dtype == bool:
+                warnings.warn(
+                    "index notation '[]' is deprecated for boolean masks. Consider using 'filter'.",
+                    DeprecationWarning,
+                )
                 return self._from_pydf(self._df.filter(pli.Series("", item).inner()))
 
         if isinstance(item, Sequence):

py-polars/polars/internals/series.py

@@ -4963,6 +4963,19 @@ def year(self) -> Series:
         """
         return wrap_s(self._s.year())
 
+    def quarter(self) -> Series:
+        """
+        Extract quarter from underlying Date representation.
+        Can be performed on Date and Datetime.
+
+        Returns the quarter ranging from 1 to 4.
+
+        Returns
+        -------
+        Quarter as UInt32
+        """
+        return pli.select(pli.lit(wrap_s(self._s)).dt.quarter()).to_series()
+
     def month(self) -> Series:
         """
         Extract the month from the underlying date representation.

py-polars/src/lazy/dsl.rs

@@ -692,6 +692,9 @@ impl PyExpr {
     pub fn year(&self) -> PyExpr {
         self.clone().inner.dt().year().into()
     }
+    pub fn quarter(&self) -> PyExpr {
+        self.clone().inner.dt().quarter().into()
+    }
     pub fn month(&self) -> PyExpr {
         self.clone().inner.dt().month().into()
     }

py-polars/tests/test_datelike.py

@@ -1149,3 +1149,9 @@ def test_groupby_rolling_by_() -> None:
         ],
         "count": [1, 2, 3, 3, 3, 1, 2, 3, 3, 3, 1, 2, 3, 3, 3],
     }
+
+
+def test_quarter() -> None:
+    assert pl.date_range(
+        datetime(2022, 1, 1), datetime(2022, 12, 1), "1mo"
+    ).dt.quarter().to_list() == [1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4]