cw_disjoint_partition_by_regexp.sqlx

config { hasOutput: true }
/*
 * Copyright 2024 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

CREATE OR REPLACE FUNCTION ${self()}(firstRn INT64, haystack STRING, regex STRING)
  RETURNS ARRAY<INT64> 
  OPTIONS(description="""Returns disjoint matches for the regex with the sequence of rows encoded in
 custom format. The expected input haystack format is
 "row-1@row-number-1#row-2@row-number-2# ... #row-n@row-number-n#".
 If the regex matches then, it returns the rows matched by the given pattern
 if the firstRn is the first row in the given matched subsequence. e.g. if
 firstRn is 4,  haystack is "A@1#A@2#B@3#A@4#B@5#" and regex is
 "(?:A@\\d+#)+(?:B@\\d+#)" then it will return [4, 5] since matched subsequence
 is "A@4#B@5#" starts at firstRn 4, however if you have input firstRn is 5 with
 the same input then it will return an empty array since row number 5 is part
 of "A@4#B@5#" but it is not the first row in the subsequence.

 By repeatedly calling cw_disjoint_partition_by_regexp with increasing values of
 firstRn and keeping haystack constant, this UDF effectively returns disjoint
 sets of row numbers that match the given regex.

 Continuing the above example, if we call this UDF with firstRn having
 increasing values, then we get corresponding outputs as follows:
 firstRn       haystack                output
 1             "A@1#A@2#B@3#A@4#B@5#"  [1, 2, 3]
 2             "A@1#A@2#B@3#A@4#B@5#"  []
 3             "A@1#A@2#B@3#A@4#B@5#"  []
 4             "A@1#A@2#B@3#A@4#B@5#"  [4, 5]
 5             "A@1#A@2#B@3#A@4#B@5#"  []
""")
  AS (
    (WITH t AS (
        SELECT
          MIN(ARRAY_LENGTH(REGEXP_EXTRACT_ALL(m, '@\\d+#'))) AS n_rows
        FROM UNNEST(REGEXP_EXTRACT_ALL(haystack, regex)) m
        WHERE REGEXP_EXTRACT(m, '@(\\d+)#') = CAST(firstRn AS STRING)
      )
      SELECT GENERATE_ARRAY(firstRn, firstRn + n_rows - 1) FROM t)
  );