apache / datafusion

Apache DataFusion SQL Query Engine
https://datafusion.apache.org/
Apache License 2.0
5.49k stars 1.02k forks source link

Implement `REGEXP_REPLACE` for StringView #11025

Open alamb opened 1 week ago

alamb commented 1 week ago

Is your feature request related to a problem or challenge?

Part of https://github.com/apache/datafusion/issues/10918 where we are integrating StringView into DataFusion, initially targeting making ClickBench queries faster

In the ClickBench queries there is a REGEXP_REPLACE function on String columns such as

SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25;

https://github.com/apache/datafusion/blob/5bfc11ba4ac4f11eaf9793c668e4a064fb697e6e/benchmarks/queries/clickbench/queries.sql#L29

Describe the solution you'd like

I would like to be able to run REGEXP_REPLACE on string view columns

Given this table:

> CREATE OR REPLACE TABLE string_views AS VALUES (arrow_cast('http://foo.com/ff', 'Utf8View')), ('http://bar.com/bs'), ('http://foo.com/42'), (NULL);
0 row(s) fetched.
Elapsed 0.011 seconds.

> select *, arrow_typeof(column1) from string_views;
+-------------------+------------------------------------+
| column1           | arrow_typeof(string_views.column1) |
+-------------------+------------------------------------+
| http://foo.com/ff | Utf8View                           |
| http://bar.com/bs | Utf8View                           |
| http://foo.com/42 | Utf8View                           |
|                   | Utf8View                           |
+-------------------+------------------------------------+
4 row(s) fetched.
Elapsed 0.006 seconds.

I would like to be able to run this function

> SELECT REGEXP_REPLACE(column1, '^https?://(?:www\\.)?([^/]+)/.*$', '\\1') AS k from string_views;
Error during planning: The regexp_replace function can only accept strings. Got Utf8View

It works fine if you cast the column first to a string:

> SELECT REGEXP_REPLACE(arrow_cast(column1, 'Utf8'), '^https?://(?:www\\.)?([^/]+)/.*$', '\\1') AS k from string_views;
+---------+
| k       |
+---------+
| foo.com |
| bar.com |
| foo.com |
|         |
+---------+
4 row(s) fetched.
Elapsed 0.012 seconds.

Describe alternatives you've considered

We could add a coercion rule to automatically cast Utf8View to Utf8 which we probably should do in general to make it easy to work with Utf8View

However that is inefficient as it will involve copying all the strings. It would be much better to actually implement REGEXP_REPLACE for Utf8View arrays directly

I am hoping we can figure out a straightforward pattern that can generate code for any string function that works well for StringArray as well as LargeStringArry and StringViewArray

Here is how the regexp replace function is implemented now:

https://github.com/apache/datafusion/blob/5bfc11ba4ac4f11eaf9793c668e4a064fb697e6e/datafusion/functions/src/regex/regexpreplace.rs#L142-L182

Additional context

Please remember to target the string-view branch in DataFusion, rather than main with your PR

XiangpengHao commented 1 week ago

take

XiangpengHao commented 1 week ago

I am hoping we can figure out a straightforward pattern that can generate code for any string function that works well for StringArray as well as LargeStringArry and StringViewArray

Agree, I have a prototype StringArrayType (https://github.com/apache/arrow-rs/pull/5931) for LIKE operations (like, ilike, contains, begins_with etc). A similar pattern could be applied here. Maybe we can directly use the StringArrayType trait here.