From 49b19c74ace2a53477cef413d3d7ca22b0e59006 Mon Sep 17 00:00:00 2001
From: Polina Osadcha <posadcha-ua@memsql.com>
Date: Tue, 15 Mar 2022 14:26:40 +0200
Subject: [PATCH] Added pushdown for findInSet string expression

Summary: Added pushdown to singestore-spark-connetor for FindInSet string expression

Test Plan: https://webapp.io/memsql/commits?query=repo%3Asinglestore-spark-connector+id%3A232

Reviewers: amakarovych-ua

Reviewed By: amakarovych-ua

Subscribers: engineering-list

JIRA Issues: PLAT-6082

Differential Revision: https://grizzly.internal.memcompute.com/D55566
---
 .../com/singlestore/spark/ExpressionGen.scala | 10 +++++++
 .../singlestore/spark/SQLPushdownTest.scala   | 29 +++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/src/main/scala/com/singlestore/spark/ExpressionGen.scala b/src/main/scala/com/singlestore/spark/ExpressionGen.scala
index 688c604a..00b32b6a 100644
--- a/src/main/scala/com/singlestore/spark/ExpressionGen.scala
+++ b/src/main/scala/com/singlestore/spark/ExpressionGen.scala
@@ -542,6 +542,16 @@ object ExpressionGen extends LazyLogging {
           if trimStr == UTF8String.fromString(" ") =>
         f("RTRIM", srcStr)
 
+      case FindInSet(expressionExtractor(left), utf8StringFoldableExtractor(right)) => {
+        val str_array    = right.toString.split(',')
+        var caseBranches = stringToJoinable("")
+        for (i <- 1 to str_array.length) {
+          caseBranches += Raw(s"WHEN '${str_array(i - 1)}'")
+          caseBranches += Raw(s"THEN '${i.toString}'")
+        }
+        block(Raw("CASE") + left + caseBranches + Raw("ELSE 0 END"))
+      }
+
       // TODO: case _: Levenshtein => None
 
       // ----------------------------------
diff --git a/src/test/scala/com/singlestore/spark/SQLPushdownTest.scala b/src/test/scala/com/singlestore/spark/SQLPushdownTest.scala
index 824be6c1..44916330 100644
--- a/src/test/scala/com/singlestore/spark/SQLPushdownTest.scala
+++ b/src/test/scala/com/singlestore/spark/SQLPushdownTest.scala
@@ -3047,6 +3047,35 @@ class SQLPushdownTest extends IntegrationSuiteBase with BeforeAndAfterEach with
       }
     }
 
+    describe("FindInSet") {
+      it("works") {
+        testQuery(
+          "select id, find_in_set(id, '82, 1,13,54,28,39,42, owns_house,120') from users")
+      }
+      it("constant example") {
+        testQuery(
+          "select id, find_in_set('39', '1,2,3990, 13,28,39,42,54,82,120') from users")
+      }
+      it("works with empty left argument") {
+        testQuery("select find_in_set('', '1,2,3') from users")
+      }
+      it("works with empty right argument") {
+        testQuery("select find_in_set(id, '') from users")
+      }
+      it("udf in the left argument") {
+        testQuery("select find_in_set(stringIdentity(critic_review), 'id') from movies",
+                  expectPartialPushdown = true)
+      }
+      it("udf in the right argument") {
+        testQuery("select find_in_set(critic_review, stringIdentity('id')) from movies",
+                  expectPartialPushdown = true)
+      }
+      it("joinable object in the right argument") {
+        testQuery("select find_in_set(critic_review, id) from movies",
+          expectPartialPushdown = true)
+      }
+    }
+
     describe("StringTrim") {
       it("works") {
         testQuery("select id, trim(first_name) from users")