SparkSQL - Extract multiple regex matches (using SQL only)

user_4882906 0
Viewed 3454

I have a dataset of SQL queries in raw text and another with a regular expression of all the possible table names:

# queries
+-----+----------------------------------------------+
| id  | query                                        |                  
+-----+----------------------------------------------+
| 1   | select * from table_a, table_b               |
| 2   | select * from table_c join table_d...        |
+-----+----------------------------------------------+

# regexp
'table_a|table_b|table_c|table_d'

And I wanted the following result:

# expected result
+-----+----------------------------------------------+
| id  | tables                                       |                  
+-----+----------------------------------------------+
| 1   | [table_a, table_b]                           |
| 2   | [table_c, table_d]                           |
+-----+----------------------------------------------+

But using the following SQL in Spark, all I get is the first match...

select
  id,
  regexp_extract(query, 'table_a|table_b|table_c|table_d') as tables
from queries

# actual result
+-----+----------------------------------------------+
| id  | tables                                       |                  
+-----+----------------------------------------------+
| 1   |  table_a                                     |
| 2   |  table_c                                     |
+-----+----------------------------------------------+

Is there any way to do this using only Spark SQL? This is the function I am using https://people.apache.org/~pwendell/spark-nightly/spark-master-docs/latest/api/sql/#regexp_extract

EDIT

I would also accept a solution that returned the following:

# alternative solution
+-----+----------------------------------------------+
| id  | tables                                       |                  
+-----+----------------------------------------------+
| 1   | table_a                                      |
| 1   | table_b                                      |
| 2   | table_c                                      |
| 2   | table_d                                      |
+-----+----------------------------------------------+

SOLUTION

@chlebek solved this below. I reformatted his SQL using CTEs for better readability:

with
split_queries as (
  select 
    id, 
    explode(split(query, ' ')) as col
  from queries
),
extracted_tables as (
  select 
    id, 
    regexp_extract(col, 'table_a|table_b|table_c|table_d', 0) as rx
  from split_queries
)
select
  id,
  collect_set(rx) as tables
from extracted_tables
where rx != ''
group by id

Bear in mind that the split(query, ' ') part of the query will split your SQL only by spaces. If you have other things such as tabs, line breaks, comments, etc., you should deal with these before or when splitting.

apache-spark-sql apache-spark
4 Answers

You can use another SQL function available in Spark called collect_list https://docs.databricks.com/spark/latest/spark-sql/language-manual/functions.html#collect_list. You can find another sample https://mungingdata.com/apache-spark/arraytype-columns/

Basically, applying to your code it should be

val df = spark.sql("select 1 id, 'select * from table_a, table_b' query" )
val df1 = spark.sql("select 2 id, 'select * from table_c join table_d' query" )

val df3 = df.union(df1)

df3.createOrReplaceTempView("tabla")

   spark.sql("""
select id, collect_list(tables) from (
        select id, explode(split(query, ' ')) as tables
        from tabla)
where tables like 'table%' group by id""").show

The output will be

+---+--------------------+
| id|collect_list(tables)|
+---+--------------------+
|  1| [table_a,, table_b]|
|  2|  [table_c, table_d]|
+---+--------------------+

Hope this helps

user_132314810

If you have only a few values to check you can achieve it using contains function instead of regexp:

val names = Seq("table_a","table_b","table_c","table_d")
def c(col: Column) = names.map(n => when(col.contains(n),n).otherwise(""))
df.select('id,array_remove(array(c('query):_*),"").as("result")).show(false)

but using regexp it will looks like below (Spark SQL API):

df.select('id,explode(split('query," ")))
.select('id,regexp_extract('col,"table_a|table_b|table_c|table_d",0).as("rx"))
.filter('rx=!="")
.groupBy('id)
.agg(collect_list('rx))

and it could be translated to below SQL query:

select id, collect_list(rx) from
    (select id, regexp_extract(col,'table_a|table_b|table_c|table_d',0) as rx from
        (select id, explode(split(query,' ')) as col from df) q1
    ) q2
where rx != '' group by id

so output will be:

+---+------------------+
| id|  collect_list(rx)|
+---+------------------+
|  1|[table_a, table_b]|
|  2|[table_c, table_d]|
+---+------------------+
user_96879100

As you are using spark-sql, you can use sql parser & it will do job for you.

def getTables(query: String): Seq[String] = {
  val logicalPlan = spark.sessionState.sqlParser.parsePlan(query)
  import org.apache.spark.sql.catalyst.analysis.UnresolvedRelation
  logicalPlan.collect { case r: UnresolvedRelation => r.tableName }
}

val query = "select * from table_1 as a left join table_2 as b on 
a.id=b.id"
scala> getTables(query).foreach(println)
table_1
table_2

You can register 'getTables' as udf & use in query

user_63595840

If you are on spark>=2.4 then you can remove exploding and collecting the same operations by using higher order functions on array and without any subqueries-

Load the test data

 val data =
      """
        |id  | query
        |1   | select * from table_a, table_b
        |2   | select * from table_c join table_d on table_c.id=table_d.id
      """.stripMargin
    val stringDS = data.split(System.lineSeparator())
      .map(_.split("\\|").map(_.replaceAll("""^[ \t]+|[ \t]+$""", "")).mkString(";"))
      .toSeq.toDS()
    val df = spark.read
      .option("sep", ";")
      .option("inferSchema", "true")
      .option("header", "true")
      .option("nullValue", "null")
      .csv(stringDS)
    df.printSchema()
    df.show(false)

    /**
      * root
      * |-- id: integer (nullable = true)
      * |-- query: string (nullable = true)
      *
      * +---+-----------------------------------------------------------+
      * |id |query                                                      |
      * +---+-----------------------------------------------------------+
      * |1  |select * from table_a, table_b                             |
      * |2  |select * from table_c join table_d on table_c.id=table_d.id|
      * +---+-----------------------------------------------------------+
      */

Extract the tables from query

  // spark >= 2.4.0
    df.createOrReplaceTempView("queries")
    spark.sql(
      """
        |select id,
        |   array_distinct(
        |   FILTER(
        |     split(query, '\\.|=|\\s+|,'), x -> x rlike 'table_a|table_b|table_c|table_d'
        |     )
        |    )as tables
        |FROM
        |   queries
      """.stripMargin)
      .show(false)

    /**
      * +---+------------------+
      * |id |tables            |
      * +---+------------------+
      * |1  |[table_a, table_b]|
      * |2  |[table_c, table_d]|
      * +---+------------------+
      */
user_47588230
Related