[calcite] 06/08: [CALCITE-4579] Piglet throws ClassCastException if Pig Latin script contains FLATTEN or STRSPLIT operators (Mahesh Kumar Behera)

[jh...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

jhyde pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/calcite.git

commit 760714d9ed4a1323d59b4bb2bf2887f030500019
Author: Mahesh Kumar Behera <ma...@apache.org>
AuthorDate: Mon Apr 12 07:47:22 2021 +0530

    [CALCITE-4579] Piglet throws ClassCastException if Pig Latin script contains FLATTEN or STRSPLIT operators (Mahesh Kumar Behera)
    
    Close apache/calcite#2396
---
 .../apache/calcite/piglet/PigRelOpInnerVisitor.java    | 18 ++++++++++++++----
 .../java/org/apache/calcite/test/PigRelOpTest.java     | 14 ++++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/piglet/src/main/java/org/apache/calcite/piglet/PigRelOpInnerVisitor.java b/piglet/src/main/java/org/apache/calcite/piglet/PigRelOpInnerVisitor.java
index 5d62668..727a187 100644
--- a/piglet/src/main/java/org/apache/calcite/piglet/PigRelOpInnerVisitor.java
+++ b/piglet/src/main/java/org/apache/calcite/piglet/PigRelOpInnerVisitor.java
@@ -23,6 +23,7 @@ import org.apache.calcite.rel.core.JoinRelType;
 import org.apache.calcite.rel.logical.LogicalValues;
 import org.apache.calcite.rel.type.RelDataType;
 import org.apache.calcite.rex.RexNode;
+import org.apache.calcite.sql.fun.SqlStdOperatorTable;
 import org.apache.calcite.sql.type.MultisetSqlType;
 import org.apache.calcite.sql.type.SqlTypeName;
 import org.apache.calcite.sql.type.SqlTypeUtil;
@@ -172,10 +173,19 @@ class PigRelOpInnerVisitor extends PigRelOpVisitor {
               && (dataType.getFieldCount() > 0 || dataType instanceof DynamicTupleRecordType)) {
         if (dataType instanceof DynamicTupleRecordType) {
           ((DynamicTupleRecordType) dataType).resize(outputFieldSchema.size());
-        }
-        for (int j = 0; j < dataType.getFieldCount(); j++) {
-          innerCols.add(builder.dot(rexNode, j));
-          fieldAlias.add(outputFieldSchema.getField(j).alias);
+          for (int j = 0; j < outputFieldSchema.size(); j++) {
+            final RelDataType scriptType = PigTypes.convertSchemaField(
+                outputFieldSchema.getField(j));
+            RexNode exp = builder.call(
+                SqlStdOperatorTable.ITEM, rexNode, builder.literal(j + 1));
+            innerCols.add(builder.getRexBuilder().makeCast(scriptType, exp));
+            fieldAlias.add(outputFieldSchema.getField(j).alias);
+          }
+        } else {
+          for (int j = 0; j < dataType.getFieldCount(); j++) {
+            innerCols.add(builder.dot(rexNode, j));
+            fieldAlias.add(outputFieldSchema.getField(j).alias);
+          }
         }
       } else {
         innerCols.add(rexNode);
diff --git a/piglet/src/test/java/org/apache/calcite/test/PigRelOpTest.java b/piglet/src/test/java/org/apache/calcite/test/PigRelOpTest.java
index 5c1a5b2..de14526 100644
--- a/piglet/src/test/java/org/apache/calcite/test/PigRelOpTest.java
+++ b/piglet/src/test/java/org/apache/calcite/test/PigRelOpTest.java
@@ -1607,4 +1607,18 @@ class PigRelOpTest extends PigRelTestBase {
         .assertResult(is(result))
         .assertSql(is(sql));
   }
+
+  @Test void testFlattenStrSplit() {
+    final String script = ""
+        + "A = LOAD 'scott.DEPT' as (DEPTNO:int, DNAME:chararray, LOC:CHARARRAY);\n"
+        + "B = FOREACH A GENERATE FLATTEN(STRSPLIT(DNAME, ',')) as NAMES;\n";
+    final String plan = ""
+        + "LogicalProject(NAMES=[CAST(ITEM(STRSPLIT(PIG_TUPLE($1, ',')), 1)):BINARY(1)])\n"
+        + "  LogicalTableScan(table=[[scott, DEPT]])\n";
+    final String sql = ""
+        + "SELECT CAST(STRSPLIT(PIG_TUPLE(DNAME, ','))[1] AS BINARY(1)) AS NAMES\n"
+        + "FROM scott.DEPT";
+    pig(script).assertRel(hasTree(plan))
+        .assertSql(is(sql));
+  }
 }