rewrite get_json_object in singular_or_list (apache#7551)

yikf · Oct 16, 2024 · 05eb6c8 · 05eb6c8
1 parent fd2805f
commit 05eb6c8
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 0 deletions.
diff --git a/.../test/scala/org/apache/gluten/execution/compatibility/GlutenClickhouseFunctionSuite.scala b/.../test/scala/org/apache/gluten/execution/compatibility/GlutenClickhouseFunctionSuite.scala
@@ -269,4 +269,30 @@ class GlutenClickhouseFunctionSuite extends GlutenClickHouseTPCHAbstractSuite {
     }
   }
 
+  test("GLUTEN-7550 get_json_object in IN") {
+    withTable("test_7550") {
+      sql("create table test_7550(a string) using parquet")
+      val insert_sql =
+        """
+          |insert into test_7550 values('{\'a\':\'1\'}'),('{\'a\':\'2\'}'),('{\'a\':\'3\'}')
+          |""".stripMargin
+      sql(insert_sql)
+      compareResultsAgainstVanillaSpark(
+        """
+          |select a, get_json_object(a, '$.a') in ('1', '2') from test_7550
+          |""".stripMargin,
+        true,
+        { _ => }
+      )
+      compareResultsAgainstVanillaSpark(
+        """
+          |select a in ('1', '2') from test_7550
+          |where get_json_object(a, '$.a') in ('1', '2')
+          |""".stripMargin,
+        true,
+        { _ => }
+      )
+    }
+  }
+
 }
diff --git a/cpp-ch/local-engine/Rewriter/ExpressionRewriter.h b/cpp-ch/local-engine/Rewriter/ExpressionRewriter.h
@@ -104,6 +104,10 @@ class GetJsonObjectFunctionWriter : public RelRewriter
                 prepareOnExpression(if_then.else_());
                 break;
             }
+            case substrait::Expression::RexTypeCase::kSingularOrList: {
+                prepareOnExpression(expr.singular_or_list().value());
+                break;
+            }
             case substrait::Expression::RexTypeCase::kScalarFunction: {
                 const auto & scalar_function_pb = expr.scalar_function();
                 auto function_signature_name_opt = parser_context->getFunctionNameInSignature(scalar_function_pb);
@@ -160,6 +164,10 @@ class GetJsonObjectFunctionWriter : public RelRewriter
                 rewriteExpression(*if_then->mutable_else_());
                 break;
             }
+            case substrait::Expression::RexTypeCase::kSingularOrList: {
+                rewriteExpression(*expr.mutable_singular_or_list()->mutable_value());
+                break;
+            }
             case substrait::Expression::RexTypeCase::kScalarFunction: {
                 auto & scalar_function_pb = *expr.mutable_scalar_function();
                 if (scalar_function_pb.arguments().empty())