Add new regex_capture scalar functions and regex_captures table function

Author: asg017

README.md

@@ -42,6 +42,28 @@ from regex_find_all(
 */
 ```
 
+**Extract capture group values by index or name**
+
+```sql
+select
+  regex_capture(captures, 0)        as entire_match,
+  regex_capture(captures, 'title')  as title,
+  regex_capture(captures, 'year')   as year
+from regex_captures(
+  regex("'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)"),
+  "'Citizen Kane' (1941), 'The Wizard of Oz' (1939), 'M' (1931)."
+);
+/*
+┌───────────────────────────┬──────────────────┬──────┐
+│       entire_match        │      title       │ year │
+├───────────────────────────┼──────────────────┼──────┤
+│ 'Citizen Kane' (1941)     │ Citizen Kane     │ 1941 │
+│ 'The Wizard of Oz' (1939) │ The Wizard of Oz │ 1939 │
+│ 'M' (1931)                │ M                │ 1931 │
+└───────────────────────────┴──────────────────┴──────┘
+*/
+```
+
 **Use RegexSets to match a string on multiple patterns in linear time**
 
 ```sql

docs.md

@@ -103,6 +103,147 @@ from regex_find_all(
 └───────┴───────┴─────┴───────────────┘
 ```
 
+<h3 name="regex_capture"><code>regex_capture(pattern, text, group)</code></h3>
+
+Returns the text of the capture group with the specific `group` index or name, or NULL otherwise. Errors if `pattern` is not legal regex. Based on [`Regex.captures()`](https://docs.rs/regex/latest/regex/struct.Regex.html#method.captures).
+
+If `group` is a number, then the N-th capture group is returned, where `0` refers to the entire match, `1` refers to the first left-most capture group in the match, `2` the second, and so on. If the provided group number "overflows', then NULL is returned.
+
+```sql
+select regex_capture(
+  "'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)",
+  "Not my favorite movie: 'Citizen Kane' (1941).",
+  0
+);
+-- "'Citizen Kane' (1941)"
+
+select regex_capture(
+  "'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)",
+  "Not my favorite movie: 'Citizen Kane' (1941).",
+  1
+);
+-- "Citizen Kane"
+
+
+select regex_capture(
+  "'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)",
+  "Not my favorite movie: 'Citizen Kane' (1941).",
+  2
+);
+-- "1941"
+
+select regex_capture(
+  "'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)",
+  "Not my favorite movie: 'Citizen Kane' (1941).",
+  3
+);
+-- NULL
+```
+
+If group is a string, then the value of the capture group with the same name is returned. If there is no matching capture group with the name, or the group was not captured, then NULL is returned.
+
+```sql
+select regex_capture(
+  "'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)",
+  "Not my favorite movie: 'Citizen Kane' (1941).",
+  'title'
+);
+-- "Citizen Kane"
+
+
+select regex_capture(
+  "'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)",
+  "Not my favorite movie: 'Citizen Kane' (1941).",
+  'year'
+);
+-- "1941"
+
+select regex_capture(
+  "'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)",
+  "Not my favorite movie: 'Citizen Kane' (1941).",
+  'not_exist'
+);
+-- NULL
+```
+
+Note that there is a version of `regex_capture()` that only have two parameters: `captures` and `group`. This can only be used with the [`regex_captures`](#regex_captures) table function, with the special `captures` column like so:
+
+```sql
+select
+  regex_capture(captures, 'title')      as title,
+  regex_capture(captures, 'year')       as year,
+  regex_capture(captures, 'not_exist')  as not_exist
+from regex_captures(
+  regex("'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)"),
+  "'Citizen Kane' (1941), 'The Wizard of Oz' (1939), 'M' (1931)."
+);
+/*
+┌──────────────────┬──────┬───────────┐
+│      title       │ year │ not_exist │
+├──────────────────┼──────┼───────────┤
+│ Citizen Kane     │ 1941 │           │
+│ The Wizard of Oz │ 1939 │           │
+│ M                │ 1931 │           │
+└──────────────────┴──────┴───────────┘
+*/
+```
+
+<h3 name="regex_captures"><code>select * from regex_captures(pattern, text)</code></h3>
+
+Returns all non-overlapping capture groups in the given text. Similar to [`regex_find_all`](#regex_find_all), but allows for extracting capture information. Must use with the [`regex_capture`](#regex_capture) function to extract capture group values. Based on [`Regex.captures_iter()`](https://docs.rs/regex/latest/regex/struct.Regex.html#method.captures_iter).
+
+The returned columns:
+
+- `rowid`: The 0-based index of the match. `0` is the entire match, `1` the first matching capture group, `2` the second, etc.
+- `captures`: A special value that's meant to be passed into [`regex_capture()`](#regex_capture). Will appear NULL through direct access.
+
+For faster results, wrap the pattern with the [`regex()`](#regex) function for caching.
+
+```sql
+select
+  rowid,
+  captures,
+  regex_capture(captures, 0)        as "0",
+  regex_capture(captures, 1)        as "1",
+  regex_capture(captures, 2)        as "2",
+  regex_capture(captures, 3)        as "3"
+from regex_captures(
+  regex("'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)"),
+  "'Citizen Kane' (1941), 'The Wizard of Oz' (1939), 'M' (1931)."
+);
+/*
+┌───────┬──────────┬───────────────────────────┬──────────────────┬──────┬───┐
+│ rowid │ captures │             0             │        1         │  2   │ 3 │
+├───────┼──────────┼───────────────────────────┼──────────────────┼──────┼───┤
+│ 0     │          │ 'Citizen Kane' (1941)     │ Citizen Kane     │ 1941 │   │
+│ 1     │          │ 'The Wizard of Oz' (1939) │ The Wizard of Oz │ 1939 │   │
+│ 2     │          │ 'M' (1931)                │ M                │ 1931 │   │
+└───────┴──────────┴───────────────────────────┴──────────────────┴──────┴───┘
+*/
+```
+
+```sql
+select
+  rowid,
+  captures,
+  regex_capture(captures, 'title')  as title,
+  regex_capture(captures, 'year')  as year,
+  regex_capture(captures, 'blah')  as blah
+from regex_captures(
+  regex("'(?P<title>[^']+)'\s+\((?P<year>\d{4})\)"),
+  "'Citizen Kane' (1941), 'The Wizard of Oz' (1939), 'M' (1931)."
+);
+/*
+┌───────┬──────────┬──────────────────┬──────┬──────┐
+│ rowid │ captures │      title       │ year │ blah │
+├───────┼──────────┼──────────────────┼──────┼──────┤
+│ 0     │          │ Citizen Kane     │ 1941 │      │
+│ 1     │          │ The Wizard of Oz │ 1939 │      │
+│ 2     │          │ M                │ 1931 │      │
+└───────┴──────────┴──────────────────┴──────┴──────┘
+*/
+```
+
 <h3 name="regex_replace"><code>regex_replace(pattern, text, replacement)</code></h3>
 
 Replace the **first** instance of `pattern` inside `text` with the given `replacement` text. Supports the [replacment string syntax](https://docs.rs/regex/latest/regex/struct.Regex.html#replacement-string-syntax). Based on [`Regex.replace()`](https://docs.rs/regex/latest/regex/struct.Regex.html#method.replace)

src/captures.rs

@@ -0,0 +1,180 @@
+use regex::{Captures, Regex};
+use sqlite_loadable::{
+    api,
+    table::{ConstraintOperator, IndexInfo, VTab, VTabArguments, VTabCursor},
+    BestIndexError, Result,
+};
+use sqlite_loadable::{prelude::*, Error};
+
+use std::{mem, os::raw::c_int};
+
+use crate::utils::{result_regex_captures, value_regex};
+
+static CREATE_SQL: &str = "CREATE TABLE x(captures, pattern hidden, contents text hidden)";
+enum Columns {
+    Captures,
+    Pattern,
+    Contents,
+}
+fn column(index: i32) -> Option<Columns> {
+    match index {
+        0 => Some(Columns::Captures),
+        1 => Some(Columns::Pattern),
+        2 => Some(Columns::Contents),
+        _ => None,
+    }
+}
+
+#[repr(C)]
+pub struct RegexCapturesTable {
+    /// must be first
+    base: sqlite3_vtab,
+}
+
+impl<'vtab> VTab<'vtab> for RegexCapturesTable {
+    type Aux = ();
+    type Cursor = RegexCapturesCursor<'vtab>;
+
+    fn connect(
+        _db: *mut sqlite3,
+        _aux: Option<&Self::Aux>,
+        _args: VTabArguments,
+    ) -> Result<(String, RegexCapturesTable)> {
+        let base: sqlite3_vtab = unsafe { mem::zeroed() };
+        let vtab = RegexCapturesTable { base };
+        // TODO db.config(VTabConfig::Innocuous)?;
+        Ok((CREATE_SQL.to_owned(), vtab))
+    }
+    fn destroy(&self) -> Result<()> {
+        Ok(())
+    }
+
+    fn best_index(&self, mut info: IndexInfo) -> core::result::Result<(), BestIndexError> {
+        let mut has_pattern = false;
+        let mut has_contents = false;
+        for mut constraint in info.constraints() {
+            match column(constraint.column_idx()) {
+                Some(Columns::Pattern) => {
+                    if constraint.usable() && constraint.op() == Some(ConstraintOperator::EQ) {
+                        constraint.set_omit(true);
+                        constraint.set_argv_index(1);
+                        has_pattern = true;
+                    } else {
+                        return Err(BestIndexError::Constraint);
+                    }
+                }
+                Some(Columns::Contents) => {
+                    if constraint.usable() && constraint.op() == Some(ConstraintOperator::EQ) {
+                        constraint.set_omit(true);
+                        constraint.set_argv_index(2);
+                        has_contents = true;
+                    } else {
+                        return Err(BestIndexError::Constraint);
+                    }
+                }
+                _ => (),
+            }
+        }
+        if !has_pattern || !has_contents {
+            return Err(BestIndexError::Error);
+        }
+        info.set_estimated_cost(100000.0);
+        info.set_estimated_rows(100000);
+        info.set_idxnum(2);
+
+        Ok(())
+    }
+
+    fn open(&mut self) -> Result<RegexCapturesCursor<'_>> {
+        Ok(RegexCapturesCursor::new())
+    }
+}
+
+#[repr(C)]
+pub struct RegexCapturesCursor<'vtab> {
+    /// Base class. Must be first
+    base: sqlite3_vtab_cursor,
+    r_clone: Option<Regex>,
+    all_captures: Option<Vec<Captures<'vtab>>>,
+    curr: usize,
+}
+impl RegexCapturesCursor<'_> {
+    fn new<'vtab>() -> RegexCapturesCursor<'vtab> {
+        let base: sqlite3_vtab_cursor = unsafe { mem::zeroed() };
+        RegexCapturesCursor {
+            base,
+            r_clone: None,
+            all_captures: None,
+            curr: 0,
+        }
+    }
+}
+
+impl VTabCursor for RegexCapturesCursor<'_> {
+    fn filter(
+        &mut self,
+        _idx_num: c_int,
+        _idx_str: Option<&str>,
+        values: &[*mut sqlite3_value],
+    ) -> Result<()> {
+        let r = value_regex(
+            values
+                .get(0)
+                .ok_or_else(|| Error::new_message("expected 1st argument as regex"))?,
+        )?;
+        let contents = api::value_text_notnull(
+            values
+                .get(1)
+                .ok_or_else(|| Error::new_message("expected 2nd argument as contents"))?,
+        )?;
+
+        let mut res = vec![];
+        for captures in r.captures_iter(contents) {
+            res.push(captures)
+        }
+        self.r_clone = Some((*r).clone());
+        Box::into_raw(r);
+        self.all_captures = Some(res);
+        self.curr = 0;
+        Ok(())
+    }
+
+    fn next(&mut self) -> Result<()> {
+        self.curr += 1;
+        Ok(())
+    }
+
+    fn eof(&self) -> bool {
+        self.all_captures
+            .as_ref()
+            .map_or(true, |m| self.curr >= m.len())
+    }
+
+    fn column(&self, context: *mut sqlite3_context, i: c_int) -> Result<()> {
+        let captures = self
+            .all_captures
+            .as_ref()
+            .ok_or_else(|| {
+                Error::new_message("sqlite-regex internal error: self.all_captures is not defined")
+            })?
+            .get(self.curr)
+            .ok_or_else(|| {
+                Error::new_message(
+                    "sqlite-regex internal error: self.curr greater than all_captures result",
+                )
+            })?;
+        match column(i) {
+            Some(Columns::Captures) => {
+                result_regex_captures(context, self.r_clone.as_ref().unwrap(), captures);
+            }
+            Some(Columns::Pattern) => (),
+            Some(Columns::Contents) => (),
+            None => (),
+        }
+        Ok(())
+    }
+
+    fn rowid(&self) -> Result<i64> {
+        Ok(self.curr as i64)
+    }
+}

src/lib.rs

@@ -1,3 +1,4 @@
+mod captures;
 mod find_all;
 mod meta;
 mod regex;
@@ -12,7 +13,10 @@ use sqlite_loadable::{
     define_scalar_function, define_table_function, errors::Result, FunctionFlags,
 };
 
-use crate::{find_all::RegexFindAllTable, meta::*, regex::*, regexset::*, split::RegexSplitTable};
+use crate::{
+    captures::RegexCapturesTable, find_all::RegexFindAllTable, meta::*, regex::*, regexset::*,
+    split::RegexSplitTable,
+};
 
 #[sqlite_entrypoint]
 pub fn sqlite3_regex_init(db: *mut sqlite3) -> Result<()> {
@@ -34,8 +38,12 @@ pub fn sqlite3_regex_init(db: *mut sqlite3) -> Result<()> {
     define_scalar_function(db, "regex_replace", 3, regex_replace, flags)?;
     define_scalar_function(db, "regex_replace_all", 3, regex_replace_all, flags)?;
 
+    define_scalar_function(db, "regex_capture", 3, regex_capture, flags)?;
+    define_scalar_function(db, "regex_capture", 2, regex_capture2, flags)?;
+
     define_table_function::<RegexFindAllTable>(db, "regex_find_all", None)?;
     define_table_function::<RegexSplitTable>(db, "regex_split", None)?;
+    define_table_function::<RegexCapturesTable>(db, "regex_captures", None)?;
 
     define_scalar_function(db, "regexset", -1, regexset, flags)?;
     define_scalar_function(db, "regexset_print", 1, regexset_print, flags)?;