polars-df 0.21.0 → 0.21.1

data/lib/polars/string_expr.rb

@@ -368,6 +368,71 @@ module Polars
     end
     alias_method :concat, :join
 
+    # Returns string values with all regular expression meta characters escaped.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"text" => ["abc", "def", nil, "abc(\\w+)"]})
+    #   df.with_columns(Polars.col("text").str.escape_regex.alias("escaped"))
+    #   # =>
+    #   # shape: (4, 2)
+    #   # ┌──────────┬──────────────┐
+    #   # │ text     ┆ escaped      │
+    #   # │ ---      ┆ ---          │
+    #   # │ str      ┆ str          │
+    #   # ╞══════════╪══════════════╡
+    #   # │ abc      ┆ abc          │
+    #   # │ def      ┆ def          │
+    #   # │ null     ┆ null         │
+    #   # │ abc(\w+) ┆ abc\(\\w\+\) │
+    #   # └──────────┴──────────────┘
+    def escape_regex
+      Utils.wrap_expr(_rbexpr.str_escape_regex)
+    end
+
+    # Returns the Unicode normal form of the string values.
+    #
+    # This uses the forms described in Unicode Standard Annex 15: <https://www.unicode.org/reports/tr15/>.
+    #
+    # @param form ['NFC', 'NFKC', 'NFD', 'NFKD']
+    #   Unicode form to use.
+    #
+    # @return [Expr]
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"text" => ["01²", "ＫＡＤＯＫＡＷＡ"]})
+    #   new = df.with_columns(
+    #     nfc: Polars.col("text").str.normalize("NFC"),
+    #     nfkc: Polars.col("text").str.normalize("NFKC")
+    #   )
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌──────────────────┬──────────────────┬──────────┐
+    #   # │ text             ┆ nfc              ┆ nfkc     │
+    #   # │ ---              ┆ ---              ┆ ---      │
+    #   # │ str              ┆ str              ┆ str      │
+    #   # ╞══════════════════╪══════════════════╪══════════╡
+    #   # │ 01²              ┆ 01²              ┆ 012      │
+    #   # │ ＫＡＤＯＫＡＷＡ ┆ ＫＡＤＯＫＡＷＡ ┆ KADOKAWA │
+    #   # └──────────────────┴──────────────────┴──────────┘
+    #
+    # @example
+    #   new.select(Polars.all.str.len_bytes)
+    #   # =>
+    #   # shape: (2, 3)
+    #   # ┌──────┬─────┬──────┐
+    #   # │ text ┆ nfc ┆ nfkc │
+    #   # │ ---  ┆ --- ┆ ---  │
+    #   # │ u32  ┆ u32 ┆ u32  │
+    #   # ╞══════╪═════╪══════╡
+    #   # │ 4    ┆ 4   ┆ 3    │
+    #   # │ 24   ┆ 24  ┆ 8    │
+    #   # └──────┴─────┴──────┘
+    def normalize(form = "NFC")
+      Utils.wrap_expr(_rbexpr.str_normalize(form))
+    end
+
     # Transform to uppercase variant.
     #
     # @return [Expr]
@@ -707,6 +772,68 @@ module Polars
       Utils.wrap_expr(_rbexpr.str_contains(pattern, literal, strict))
     end
 
+    # Return the bytes offset of the first substring matching a pattern.
+    #
+    # If the pattern is not found, returns None.
+    #
+    # @param pattern [String]
+    #   A valid regular expression pattern, compatible with the [regex crate](https://docs.rs/regex/latest/regex/).
+    # @param literal [Boolean]
+    #   Treat `pattern` as a literal string, not as a regular expression.
+    # @param strict [Boolean]
+    #   Raise an error if the underlying pattern is not a valid regex,
+    #   otherwise mask out with a null value.
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   To modify regular expression behaviour (such as case-sensitivity) with
+    #   flags, use the inline `(?iLmsuxU)` syntax.
+    #
+    # @example Find the index of the first substring matching a regex or literal pattern:
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "txt" => ["Crab", "Lobster", nil, "Crustacean"],
+    #       "pat" => ["a[bc]", "b.t", "[aeiuo]", "(?i)A[BC]"]
+    #     }
+    #   )
+    #   df.select(
+    #     Polars.col("txt"),
+    #     Polars.col("txt").str.find("a|e").alias("a|e (regex)"),
+    #     Polars.col("txt").str.find("e", literal: true).alias("e (lit)"),
+    #   )
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌────────────┬─────────────┬─────────┐
+    #   # │ txt        ┆ a|e (regex) ┆ e (lit) │
+    #   # │ ---        ┆ ---         ┆ ---     │
+    #   # │ str        ┆ u32         ┆ u32     │
+    #   # ╞════════════╪═════════════╪═════════╡
+    #   # │ Crab       ┆ 2           ┆ null    │
+    #   # │ Lobster    ┆ 5           ┆ 5       │
+    #   # │ null       ┆ null        ┆ null    │
+    #   # │ Crustacean ┆ 5           ┆ 7       │
+    #   # └────────────┴─────────────┴─────────┘
+    #
+    # @example Match against a pattern found in another column or (expression):
+    #   df.with_columns(Polars.col("txt").str.find(Polars.col("pat")).alias("find_pat"))
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌────────────┬───────────┬──────────┐
+    #   # │ txt        ┆ pat       ┆ find_pat │
+    #   # │ ---        ┆ ---       ┆ ---      │
+    #   # │ str        ┆ str       ┆ u32      │
+    #   # ╞════════════╪═══════════╪══════════╡
+    #   # │ Crab       ┆ a[bc]     ┆ 2        │
+    #   # │ Lobster    ┆ b.t       ┆ 2        │
+    #   # │ null       ┆ [aeiuo]   ┆ null     │
+    #   # │ Crustacean ┆ (?i)A[BC] ┆ 5        │
+    #   # └────────────┴───────────┴──────────┘
+    def find(pattern, literal: false, strict: true)
+      pattern = Utils.parse_into_expression(pattern, str_as_lit: true)
+      Utils.wrap_expr(_rbexpr.str_find(pattern, literal, strict))
+    end
+
     # Check if string values end with a substring.
     #
     # @param sub [String]
@@ -1307,6 +1434,130 @@ module Polars
       Utils.wrap_expr(_rbexpr.str_slice(offset, length))
     end
 
+    # Return the first n characters of each string in a String Series.
+    #
+    # @param n [Integer]
+    #   Length of the slice (integer or expression). Negative indexing is supported;
+    #   see note (2) below.
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   1) The `n` input is defined in terms of the number of characters in the (UTF8)
+    #      string. A character is defined as a [Unicode scalar value](https://www.unicode.org/glossary/#unicode_scalar_value). A single
+    #      character is represented by a single byte when working with ASCII text, and a
+    #      maximum of 4 bytes otherwise.
+    #
+    #   2) When the `n` input is negative, `head` returns characters up to the `n`th
+    #      from the end of the string. For example, if `n = -3`, then all characters
+    #      except the last three are returned.
+    #
+    #   3) If the length of the string has fewer than `n` characters, the full string is
+    #      returned.
+    #
+    # @example Return up to the first 5 characters:
+    #   df = Polars::DataFrame.new({"s" => ["pear", nil, "papaya", "dragonfruit"]})
+    #   df.with_columns(Polars.col("s").str.head(5).alias("s_head_5"))
+    #   # =>
+    #   # shape: (4, 2)
+    #   # ┌─────────────┬──────────┐
+    #   # │ s           ┆ s_head_5 │
+    #   # │ ---         ┆ ---      │
+    #   # │ str         ┆ str      │
+    #   # ╞═════════════╪══════════╡
+    #   # │ pear        ┆ pear     │
+    #   # │ null        ┆ null     │
+    #   # │ papaya      ┆ papay    │
+    #   # │ dragonfruit ┆ drago    │
+    #   # └─────────────┴──────────┘
+    #
+    # @example Return characters determined by column `n`:
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "s" => ["pear", nil, "papaya", "dragonfruit"],
+    #       "n" => [3, 4, -2, -5]
+    #     }
+    #   )
+    #   df.with_columns(Polars.col("s").str.head("n").alias("s_head_n"))
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌─────────────┬─────┬──────────┐
+    #   # │ s           ┆ n   ┆ s_head_n │
+    #   # │ ---         ┆ --- ┆ ---      │
+    #   # │ str         ┆ i64 ┆ str      │
+    #   # ╞═════════════╪═════╪══════════╡
+    #   # │ pear        ┆ 3   ┆ pea      │
+    #   # │ null        ┆ 4   ┆ null     │
+    #   # │ papaya      ┆ -2  ┆ papa     │
+    #   # │ dragonfruit ┆ -5  ┆ dragon   │
+    #   # └─────────────┴─────┴──────────┘
+    def head(n)
+      n = Utils.parse_into_expression(n)
+      Utils.wrap_expr(_rbexpr.str_head(n))
+    end
+
+    # Return the last n characters of each string in a String Series.
+    #
+    # @param n [Integer]
+    #   Length of the slice (integer or expression). Negative indexing is supported;
+    #   see note (2) below.
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   1) The `n` input is defined in terms of the number of characters in the (UTF8)
+    #      string. A character is defined as a [Unicode scalar value](https://www.unicode.org/glossary/#unicode_scalar_value). A single
+    #      character is represented by a single byte when working with ASCII text, and a
+    #      maximum of 4 bytes otherwise.
+    #
+    #   2) When the `n` input is negative, `tail` returns characters starting from the
+    #      `n`th from the beginning of the string. For example, if `n = -3`, then all
+    #      characters except the first three are returned.
+    #
+    #   3) If the length of the string has fewer than `n` characters, the full string is
+    #      returned.
+    #
+    # @example Return up to the last 5 characters:
+    #   df = Polars::DataFrame.new({"s" => ["pear", nil, "papaya", "dragonfruit"]})
+    #   df.with_columns(Polars.col("s").str.tail(5).alias("s_tail_5"))
+    #   # =>
+    #   # shape: (4, 2)
+    #   # ┌─────────────┬──────────┐
+    #   # │ s           ┆ s_tail_5 │
+    #   # │ ---         ┆ ---      │
+    #   # │ str         ┆ str      │
+    #   # ╞═════════════╪══════════╡
+    #   # │ pear        ┆ pear     │
+    #   # │ null        ┆ null     │
+    #   # │ papaya      ┆ apaya    │
+    #   # │ dragonfruit ┆ fruit    │
+    #   # └─────────────┴──────────┘
+    #
+    # @example Return characters determined by column `n`:
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "s" => ["pear", nil, "papaya", "dragonfruit"],
+    #       "n" => [3, 4, -2, -5]
+    #     }
+    #   )
+    #   df.with_columns(Polars.col("s").str.tail("n").alias("s_tail_n"))
+    #   # =>
+    #   # shape: (4, 3)
+    #   # ┌─────────────┬─────┬──────────┐
+    #   # │ s           ┆ n   ┆ s_tail_n │
+    #   # │ ---         ┆ --- ┆ ---      │
+    #   # │ str         ┆ i64 ┆ str      │
+    #   # ╞═════════════╪═════╪══════════╡
+    #   # │ pear        ┆ 3   ┆ ear      │
+    #   # │ null        ┆ 4   ┆ null     │
+    #   # │ papaya      ┆ -2  ┆ paya     │
+    #   # │ dragonfruit ┆ -5  ┆ nfruit   │
+    #   # └─────────────┴─────┴──────────┘
+    def tail(n)
+      n = Utils.parse_into_expression(n)
+      Utils.wrap_expr(_rbexpr.str_tail(n))
+    end
+
     # Convert an Utf8 column into an Int64 column with base radix.
     #
     # @param base [Integer]
@@ -1432,9 +1683,9 @@ module Polars
 
     # Use the aho-corasick algorithm to replace many matches.
     #
-    # @param patterns [String]
+    # @param patterns [Object]
     #   String patterns to search and replace.
-    # @param replace_with [String]
+    # @param replace_with [Object]
     #   Strings to replace where a pattern was a match.
     #   This can be broadcasted. So it supports many:one and many:many.
     # @param ascii_case_insensitive [Boolean]
@@ -1494,11 +1745,22 @@ module Polars
     #   # │ Tell me what you want, what yo… ┆ Tell you what me want, what me… │
     #   # │ Can you feel the love tonight   ┆ Can me feel the love tonight    │
     #   # └─────────────────────────────────┴─────────────────────────────────┘
-    def replace_many(patterns, replace_with, ascii_case_insensitive: false)
+    def replace_many(patterns, replace_with = Expr::NO_DEFAULT, ascii_case_insensitive: false)
+      if replace_with == Expr::NO_DEFAULT
+        if !patterns.is_a?(Hash)
+          msg = "`replace_with` argument is required if `patterns` argument is not a Hash type"
+          raise TypeError, msg
+        end
+        # Early return in case of an empty mapping.
+        if patterns.empty?
+          return Utils.wrap_expr(_rbexpr)
+        end
+        replace_with = patterns.values
+        patterns = patterns.keys
+      end
+
       patterns = Utils.parse_into_expression(patterns, str_as_lit: false)
-      replace_with = Utils.parse_into_expression(
-        replace_with, str_as_lit: true
-      )
+      replace_with = Utils.parse_into_expression(replace_with, str_as_lit: true)
       Utils.wrap_expr(
         _rbexpr.str_replace_many(
           patterns, replace_with, ascii_case_insensitive
@@ -1506,6 +1768,149 @@ module Polars
       )
     end
 
+    # Use the Aho-Corasick algorithm to extract many matches.
+    #
+    # @param patterns [Object]
+    #   String patterns to search.
+    # @param ascii_case_insensitive [Boolean]
+    #   Enable ASCII-aware case-insensitive matching.
+    #   When this option is enabled, searching will be performed without respect
+    #   to case for ASCII letters (a-z and A-Z) only.
+    # @param overlapping [Boolean]
+    #   Whether matches may overlap.
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   This method supports matching on string literals only, and does not support
+    #   regular expression matching.
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"values" => ["discontent"]})
+    #   patterns = ["winter", "disco", "onte", "discontent"]
+    #   df.with_columns(
+    #     Polars.col("values")
+    #     .str.extract_many(patterns, overlapping: false)
+    #     .alias("matches"),
+    #     Polars.col("values")
+    #     .str.extract_many(patterns, overlapping: true)
+    #     .alias("matches_overlapping"),
+    #   )
+    #   # =>
+    #   # shape: (1, 3)
+    #   # ┌────────────┬───────────┬─────────────────────────────────┐
+    #   # │ values     ┆ matches   ┆ matches_overlapping             │
+    #   # │ ---        ┆ ---       ┆ ---                             │
+    #   # │ str        ┆ list[str] ┆ list[str]                       │
+    #   # ╞════════════╪═══════════╪═════════════════════════════════╡
+    #   # │ discontent ┆ ["disco"] ┆ ["disco", "onte", "discontent"… │
+    #   # └────────────┴───────────┴─────────────────────────────────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "values" => ["discontent", "rhapsody"],
+    #       "patterns" => [
+    #         ["winter", "disco", "onte", "discontent"],
+    #         ["rhap", "ody", "coalesce"]
+    #       ]
+    #     }
+    #   )
+    #   df.select(Polars.col("values").str.extract_many("patterns"))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌─────────────────┐
+    #   # │ values          │
+    #   # │ ---             │
+    #   # │ list[str]       │
+    #   # ╞═════════════════╡
+    #   # │ ["disco"]       │
+    #   # │ ["rhap", "ody"] │
+    #   # └─────────────────┘
+    def extract_many(
+      patterns,
+      ascii_case_insensitive: false,
+      overlapping: false
+    )
+      patterns = Utils.parse_into_expression(patterns, str_as_lit: false)
+      Utils.wrap_expr(
+        _rbexpr.str_extract_many(patterns, ascii_case_insensitive, overlapping)
+      )
+    end
+
+    # Use the Aho-Corasick algorithm to find many matches.
+    #
+    # The function will return the bytes offset of the start of each match.
+    # The return type will be `List<UInt32>`
+    #
+    # @param patterns [Object]
+    #   String patterns to search.
+    # @param ascii_case_insensitive [Boolean]
+    #   Enable ASCII-aware case-insensitive matching.
+    #   When this option is enabled, searching will be performed without respect
+    #   to case for ASCII letters (a-z and A-Z) only.
+    # @param overlapping [Boolean]
+    #   Whether matches may overlap.
+    #
+    # @return [Expr]
+    #
+    # @note
+    #   This method supports matching on string literals only, and does not support
+    #   regular expression matching.
+    #
+    # @example
+    #   df = Polars::DataFrame.new({"values" => ["discontent"]})
+    #   patterns = ["winter", "disco", "onte", "discontent"]
+    #   df.with_columns(
+    #     Polars.col("values")
+    #     .str.extract_many(patterns, overlapping: false)
+    #     .alias("matches"),
+    #     Polars.col("values")
+    #     .str.extract_many(patterns, overlapping: true)
+    #     .alias("matches_overlapping"),
+    #   )
+    #   # =>
+    #   # shape: (1, 3)
+    #   # ┌────────────┬───────────┬─────────────────────────────────┐
+    #   # │ values     ┆ matches   ┆ matches_overlapping             │
+    #   # │ ---        ┆ ---       ┆ ---                             │
+    #   # │ str        ┆ list[str] ┆ list[str]                       │
+    #   # ╞════════════╪═══════════╪═════════════════════════════════╡
+    #   # │ discontent ┆ ["disco"] ┆ ["disco", "onte", "discontent"… │
+    #   # └────────────┴───────────┴─────────────────────────────────┘
+    #
+    # @example
+    #   df = Polars::DataFrame.new(
+    #     {
+    #       "values" => ["discontent", "rhapsody"],
+    #       "patterns" => [
+    #         ["winter", "disco", "onte", "discontent"],
+    #         ["rhap", "ody", "coalesce"]
+    #       ]
+    #     }
+    #   )
+    #   df.select(Polars.col("values").str.find_many("patterns"))
+    #   # =>
+    #   # shape: (2, 1)
+    #   # ┌───────────┐
+    #   # │ values    │
+    #   # │ ---       │
+    #   # │ list[u32] │
+    #   # ╞═══════════╡
+    #   # │ [0]       │
+    #   # │ [0, 5]    │
+    #   # └───────────┘
+    def find_many(
+      patterns,
+      ascii_case_insensitive: false,
+      overlapping: false
+    )
+      patterns = Utils.parse_into_expression(patterns, str_as_lit: false)
+      Utils.wrap_expr(
+        _rbexpr.str_find_many(patterns, ascii_case_insensitive, overlapping)
+      )
+    end
+
     private
 
     def _validate_format_argument(format)