From a4d910f1a5311cb166795e37372bc955f917ac48 Mon Sep 17 00:00:00 2001
From: Alexander Beedie <alexander-beedie@users.noreply.github.com>
Date: Wed, 4 Oct 2023 16:32:09 +0400
Subject: [PATCH] feat(python): add `INITCAP` string function for SQL (#9884)

---
 crates/polars-sql/src/functions.rs   |  7 +++++++
 py-polars/tests/unit/sql/test_sql.py | 23 +++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/crates/polars-sql/src/functions.rs b/crates/polars-sql/src/functions.rs
index f39394598b377..f34289937aee9 100644
--- a/crates/polars-sql/src/functions.rs
+++ b/crates/polars-sql/src/functions.rs
@@ -182,6 +182,11 @@ pub(crate) enum PolarsSqlFunctions {
     /// SELECT column_2 from df WHERE ENDS_WITH(column_1, 'a');
     /// ```
     EndsWith,
+    /// SQL 'initcap' function
+    /// ```sql
+    /// SELECT INITCAP(column_1) from df;
+    /// ```
+    InitCap,
     /// SQL 'left' function
     /// ```sql
     /// SELECT LEFT(column_1, 3) from df;
@@ -499,6 +504,7 @@ impl PolarsSqlFunctions {
             // String functions
             // ----
             "ends_with" => Self::EndsWith,
+            "initcap" => Self::InitCap,
             "length" => Self::Length,
             "left" => Self::Left,
             "lower" => Self::Lower,
@@ -616,6 +622,7 @@ impl SqlFunctionVisitor<'_> {
             // String functions
             // ----
             EndsWith => self.visit_binary(|e, s| e.str().ends_with(s)),
+            InitCap => self.visit_unary(|e| e.str().to_titlecase()),
             Left => self.try_visit_binary(|e, length| {
                 Ok(e.str().str_slice(0, match length {
                     Expr::Literal(LiteralValue::Int64(n)) => Some(n as u64),
diff --git a/py-polars/tests/unit/sql/test_sql.py b/py-polars/tests/unit/sql/test_sql.py
index ca29917029a5d..4d5cefb70adbc 100644
--- a/py-polars/tests/unit/sql/test_sql.py
+++ b/py-polars/tests/unit/sql/test_sql.py
@@ -809,6 +809,29 @@ def test_sql_round_ndigits_errors() -> None:
         ctx.execute("SELECT ROUND(n,-1) AS n FROM df")
 
 
+def test_sql_string_case() -> None:
+    df = pl.DataFrame({"words": ["Test SOME words"]})
+
+    with pl.SQLContext(frame=df) as ctx:
+        res = ctx.execute(
+            """
+            SELECT
+              words,
+              INITCAP(words) as cap,
+              UPPER(words) as upper,
+              LOWER(words) as lower,
+            FROM frame
+            """
+        ).collect()
+
+        assert res.to_dict(False) == {
+            "words": ["Test SOME words"],
+            "cap": ["Test Some Words"],
+            "upper": ["TEST SOME WORDS"],
+            "lower": ["test some words"],
+        }
+
+
 def test_sql_string_lengths() -> None:
     df = pl.DataFrame({"words": ["Café", None, "東京"]})