Skip to content

Commit 09b4304

Browse files
acarbonettoacarbonetto
committed
Add IT tests for json_set, json_expand
Signed-off-by: Andrew Carbonetto <andrew.carbonetto@improving.com>
1 parent f20b0ee commit 09b4304

File tree

1 file changed

+257
-2
lines changed

1 file changed

+257
-2
lines changed

‎integ-test/src/integration/scala/org/opensearch/flint/spark/ppl/FlintSparkPPLJsonFunctionITSuite.scala‎

Lines changed: 257 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -507,6 +507,97 @@ class FlintSparkPPLJsonFunctionITSuite
507507
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
508508
}
509509

510+
// JSON_SET
511+
512+
test("test json_set() function: one key") {
513+
val frame = sql(s"""
514+
| source = $testTable
515+
| | eval result = json_set('$validJson1',array('age', '42')) | head 1 | fields result
516+
| """.stripMargin)
517+
assertSameRows(Seq(Row("{\"account_number\":1,\"balance\":39225,\"age\",'42',\"gender\":\"M\"}")), frame)
518+
519+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
520+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
521+
val keysExpression = UnresolvedFunction("array", Seq(Literal("age"), Literal("42")), isDistinct = false)
522+
val jsonObjExp =
523+
Literal("{\"account_number\":1,\"balance\":39225,\"age\":32,\"gender\":\"M\"}")
524+
val jsonFunc =
525+
Alias(visit("json_set", util.List.of(jsonObjExp, keysExpression)), "result")()
526+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
527+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
528+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
529+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
530+
}
531+
532+
test("test json_set() function: multiple keys") {
533+
val frame = sql(s"""
534+
| source = $testTable
535+
| | eval result = json_set('$validJson1',array('age','42','name','\"Foobar\"')) | head 1 | fields result
536+
| """.stripMargin)
537+
assertSameRows(Seq(Row("{\"account_number\":1,\"balance\":39225,\"age\",'42',\"gender\":\"M\",\"name\":\"Foobar\"}")), frame)
538+
539+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
540+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
541+
val keysExpression =
542+
UnresolvedFunction("array", Seq(Literal("age"), Literal("42"), Literal("name"), Literal("\"Foobar\"")), isDistinct = false)
543+
val jsonObjExp =
544+
Literal("{\"account_number\":1,\"balance\":39225,\"age\":32,\"gender\":\"M\"}")
545+
val jsonFunc =
546+
Alias(visit("json_set", util.List.of(jsonObjExp, keysExpression)), "result")()
547+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
548+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
549+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
550+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
551+
}
552+
553+
test("test json_set() function: nested key") {
554+
val frame = sql(s"""
555+
| source = $testTable
556+
| | eval result = json_set('$validJson2',array('f2.f3','\"zzz\"')) | head 1 | fields result
557+
| """.stripMargin)
558+
assertSameRows(Seq(Row("{\"f1\":\"abc\",\"f2\":{\"f3\":\"zzz\",\"f4\":\"b\"}}")), frame)
559+
560+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
561+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
562+
val keysExpression =
563+
UnresolvedFunction("array", Seq(Literal("f2.f3", Literal("\"zzz\""))), isDistinct = false)
564+
val jsonObjExp =
565+
Literal("{\"f1\":\"abc\",\"f2\":{\"f3\":\"a\",\"f4\":\"b\"}}")
566+
val jsonFunc =
567+
Alias(visit("json_set", util.List.of(jsonObjExp, keysExpression)), "result")()
568+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
569+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
570+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
571+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
572+
}
573+
574+
test("test json_set() function: set new key") {
575+
val frame = sql(s"""
576+
| source = $testTable
577+
| | eval result = json_set('$validJson5',array('grade','8')) | head 1 | fields result
578+
| """.stripMargin)
579+
assertSameRows(
580+
Seq(Row(
581+
"{\"teacher\":\"Alice\",\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}],\"grade\":8}")),
582+
frame)
583+
584+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
585+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
586+
val keysExpression =
587+
UnresolvedFunction("array", Seq(Literal("grade"), Literal(8)), isDistinct = false)
588+
val jsonObjExp =
589+
Literal(
590+
"{\"teacher\":\"Alice\",\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")
591+
val jsonFunc =
592+
Alias(visit("json_set", util.List.of(jsonObjExp, keysExpression)), "result")()
593+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
594+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
595+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
596+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
597+
}
598+
599+
// JSON_APPEND
600+
510601
test("test json_append() function: add single value") {
511602
val frame = sql(s"""
512603
| source = $testTable
@@ -616,7 +707,7 @@ class FlintSparkPPLJsonFunctionITSuite
616707
test("test json_append() function: add multi value") {
617708
val frame = sql(s"""
618709
| source = $testTable
619-
| | eval result = json_append('$validJson7',array('teacher', 'Tom', 'Walt')) | head 1 | fields result
710+
| | eval result = json_append('$validJson7',array('teacher', 'Tom', 'teacher' 'Walt')) | head 1 | fields result
620711
| """.stripMargin)
621712
assertSameRows(
622713
Seq(Row(
@@ -644,7 +735,7 @@ class FlintSparkPPLJsonFunctionITSuite
644735
test("test json_append() function: add nested value") {
645736
val frame = sql(s"""
646737
| source = $testTable
647-
| | eval result = json_append('$validJson8',array('school.teacher', 'Tom', 'Walt')) | head 1 | fields result
738+
| | eval result = json_append('$validJson8',array('school.teacher', 'Tom', 'school.teacher', 'Walt')) | head 1 | fields result
648739
| """.stripMargin)
649740
assertSameRows(
650741
Seq(Row(
@@ -668,4 +759,168 @@ class FlintSparkPPLJsonFunctionITSuite
668759
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
669760
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
670761
}
762+
763+
// JSON_EXTEND
764+
765+
test("test json_extend() function: add single value") {
766+
val frame = sql(s"""
767+
| source = $testTable
768+
| | eval result = json_extend('$validJson7',array('teacher', 'Tom')) | head 1 | fields result
769+
| """.stripMargin)
770+
assertSameRows(
771+
Seq(Row(
772+
"{\"teacher\":[\"Alice\",\"Tom\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")),
773+
frame)
774+
775+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
776+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
777+
val keysExpression =
778+
UnresolvedFunction("array", Seq(Literal("teacher"), Literal("Tom")), isDistinct = false)
779+
val jsonObjExp =
780+
Literal(
781+
"{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")
782+
val jsonFunc =
783+
Alias(visit("json_extend", util.List.of(jsonObjExp, keysExpression)), "result")()
784+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
785+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
786+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
787+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
788+
}
789+
790+
test("test json_extend() function: add single value key not found") {
791+
val frame = sql(s"""
792+
| source = $testTable
793+
| | eval result = json_extend('$validJson7',array('headmaster', 'Tom')) | head 1 | fields result
794+
| """.stripMargin)
795+
assertSameRows(
796+
Seq(Row(
797+
"{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}],\"headmaster\":[\"Tom\"]}")),
798+
frame)
799+
800+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
801+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
802+
val keysExpression =
803+
UnresolvedFunction("array", Seq(Literal("headmaster"), Literal("Tom")), isDistinct = false)
804+
val jsonObjExp =
805+
Literal(
806+
"{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")
807+
val jsonFunc =
808+
Alias(visit("json_extend", util.List.of(jsonObjExp, keysExpression)), "result")()
809+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
810+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
811+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
812+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
813+
}
814+
815+
test("test json_extend() function: add single Object key not found") {
816+
val frame = sql(s"""
817+
| source = $testTable
818+
| | eval result = json_extend('$validJson7',array('headmaster', '{"name":"Tomy","rank":1}')) | head 1 | fields result
819+
| """.stripMargin)
820+
assertSameRows(
821+
Seq(Row(
822+
"{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}],\"headmaster\":[{\"name\":\"Tomy\",\"rank\":1}]}")),
823+
frame)
824+
825+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
826+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
827+
val keysExpression =
828+
UnresolvedFunction(
829+
"array",
830+
Seq(Literal("headmaster"), Literal("""{"name":"Tomy","rank":1}""")),
831+
isDistinct = false)
832+
val jsonObjExp =
833+
Literal(
834+
"{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")
835+
val jsonFunc =
836+
Alias(visit("json_extend", util.List.of(jsonObjExp, keysExpression)), "result")()
837+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
838+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
839+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
840+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
841+
}
842+
843+
test("test json_extend() function: add single Object value") {
844+
val frame = sql(s"""
845+
| source = $testTable
846+
| | eval result = json_extend('$validJson7',array('student', '{"name":"Tomy","rank":5}')) | head 1 | fields result
847+
| """.stripMargin)
848+
assertSameRows(
849+
Seq(Row(
850+
"{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2},{\"name\":\"Tomy\",\"rank\":5}]}")),
851+
frame)
852+
853+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
854+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
855+
val keysExpression =
856+
UnresolvedFunction(
857+
"array",
858+
Seq(Literal("student"), Literal("""{"name":"Tomy","rank":5}""")),
859+
isDistinct = false)
860+
val jsonObjExp =
861+
Literal(
862+
"{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")
863+
val jsonFunc =
864+
Alias(visit("json_extend", util.List.of(jsonObjExp, keysExpression)), "result")()
865+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
866+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
867+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
868+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
869+
}
870+
871+
test("test json_extend() function: add multi value") {
872+
val frame = sql(s"""
873+
| source = $testTable
874+
| | eval result = json_extend('$validJson7',array('teacher', array('Tom', 'Walt'))) | head 1 | fields result
875+
| """.stripMargin)
876+
assertSameRows(
877+
Seq(Row(
878+
"{\"teacher\":[\"Alice\",\"Tom\",\"Walt\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")),
879+
frame)
880+
881+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
882+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
883+
val keysExpression =
884+
UnresolvedFunction(
885+
"array",
886+
Seq(Literal("teacher"), Literal("Tom"), Literal("Walt")),
887+
isDistinct = false)
888+
val jsonObjExp =
889+
Literal(
890+
"{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}")
891+
val jsonFunc =
892+
Alias(visit("json_extend", util.List.of(jsonObjExp, keysExpression)), "result")()
893+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
894+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
895+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
896+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
897+
}
898+
899+
test("test json_extend() function: add nested value") {
900+
val frame = sql(s"""
901+
| source = $testTable
902+
| | eval result = json_extend('$validJson8',array('school.teacher', array('Tom', 'Walt'))) | head 1 | fields result
903+
| """.stripMargin)
904+
assertSameRows(
905+
Seq(Row(
906+
"{\"school\":{\"teacher\":[\"Alice\",\"Tom\",\"Walt\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}}")),
907+
frame)
908+
909+
val logicalPlan: LogicalPlan = frame.queryExecution.logical
910+
val table = UnresolvedRelation(Seq("spark_catalog", "default", "flint_ppl_test"))
911+
val keysExpression =
912+
UnresolvedFunction(
913+
"array",
914+
Seq(Literal("school.teacher"), Literal("Tom"), Literal("Walt")),
915+
isDistinct = false)
916+
val jsonObjExp =
917+
Literal(
918+
"{\"school\":{\"teacher\":[\"Alice\"],\"student\":[{\"name\":\"Bob\",\"rank\":1},{\"name\":\"Charlie\",\"rank\":2}]}}")
919+
val jsonFunc =
920+
Alias(visit("json_extend", util.List.of(jsonObjExp, keysExpression)), "result")()
921+
val eval = Project(Seq(UnresolvedStar(None), jsonFunc), table)
922+
val limit = GlobalLimit(Literal(1), LocalLimit(Literal(1), eval))
923+
val expectedPlan = Project(Seq(UnresolvedAttribute("result")), limit)
924+
comparePlans(logicalPlan, expectedPlan, checkAnalysis = false)
925+
}
671926
}

0 commit comments

Comments
 (0)