[CALCITE-4571] In piglet, a Pig Latin script with multiple STORE comm…

…ands causes the merging of multiple SQL statements (Mahesh Kumar Behera) Close apache#2390
liyubin117 · Apr 19, 2021 · 8169246 · 8169246
1 parent 296b84c
commit 8169246
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 0 deletions.
diff --git a/piglet/src/main/java/org/apache/calcite/piglet/PigConverter.java b/piglet/src/main/java/org/apache/calcite/piglet/PigConverter.java
@@ -254,6 +254,7 @@ private List<String> pigToSql(String pigQuery, SqlWriter writer)
       final SqlNode sqlNode = sqlConverter.visitRoot(rel).asStatement();
       sqlNode.unparse(writer, 0, 0);
       sqlStatements.add(writer.toString());
+      writer.reset();
     }
     return sqlStatements;
   }

diff --git a/piglet/src/test/java/org/apache/calcite/test/PigRelOpTest.java b/piglet/src/test/java/org/apache/calcite/test/PigRelOpTest.java
@@ -105,6 +105,17 @@ private Fluent assertSql(Matcher<String> sqlMatcher) {
       }
     }
 
+    private Fluent assertSql(Matcher<String> sqlMatcher, int pos) {
+      try {
+        final String sql =
+            converter.pigToSql(script, PigRelSqlDialect.DEFAULT).get(pos);
+        assertThat(sql, sqlMatcher);
+        return this;
+      } catch (IOException e) {
+        throw TestUtil.rethrow(e);
+      }
+    }
+
     private Fluent assertResult(Matcher<String> resultMatcher) {
       final RelNode rel;
       try {
@@ -1622,6 +1633,29 @@ private Fluent pig(String script) {
         .assertSql(is(sql));
   }
 
+  @Test void testMultipleStores() {
+    final String script = ""
+        + "A = LOAD 'scott.DEPT' as (DEPTNO:int, DNAME:chararray, LOC:CHARARRAY);\n"
+        + "B = FILTER A BY DEPTNO <= 30;\n"
+        + "STORE B into 'output.csv';\n"
+        + "C = FILTER A BY DEPTNO >= 20;\n"
+        + "STORE C into 'output1.csv';\n";
+    final String plan = ""
+        + "LogicalFilter(condition=[<=($0, 30)])\n"
+        + "  LogicalTableScan(table=[[scott, DEPT]])\n";
+    final String sql0 = ""
+        + "SELECT *\n"
+        + "FROM scott.DEPT\n"
+        + "WHERE DEPTNO <= 30";
+    final String sql1 = ""
+        + "SELECT *\n"
+        + "FROM scott.DEPT\n"
+        + "WHERE DEPTNO >= 20";
+    pig(script).assertRel(hasTree(plan))
+        .assertSql(is(sql0), 0)
+        .assertSql(is(sql1), 1);
+  }
+
   @Test void testRankAndFilter() {
     final String script = ""
         + "A = LOAD 'emp1' USING PigStorage(',')  as ("
@@ -1643,4 +1677,5 @@ private Fluent pig(String script) {
     pig(script).assertRel(hasTree(plan))
         .assertSql(is(sql));
   }
+
 }