Support json type in PinotSegmentPageSource

prestodb · Feb 18, 2022 · 6b5da9a · 6b5da9a
1 parent 4d8aad1
commit 6b5da9a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotSegmentPageSource.java b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/PinotSegmentPageSource.java
@@ -494,6 +494,7 @@ private Slice getSlice(int rowIndex, int columnIndex)
                 String[] stringArray = currentDataTable.getDataTable().getStringArray(rowIndex, columnIndex);
                 return utf8Slice(Arrays.toString(stringArray));
             case STRING:
+            case JSON:
                 String field = currentDataTable.getDataTable().getString(rowIndex, columnIndex);
                 if (field == null || field.isEmpty()) {
                     return Slices.EMPTY_SLICE;

diff --git a/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/query/PinotQueryGenerator.java b/presto-pinot-toolkit/src/main/java/com/facebook/presto/pinot/query/PinotQueryGenerator.java
@@ -16,6 +16,7 @@
 import com.facebook.airlift.log.Logger;
 import com.facebook.presto.common.type.BigintType;
 import com.facebook.presto.common.type.FixedWidthType;
+import com.facebook.presto.common.type.JsonType;
 import com.facebook.presto.common.type.TypeManager;
 import com.facebook.presto.common.type.VarcharType;
 import com.facebook.presto.pinot.PinotColumnHandle;
@@ -499,7 +500,7 @@ public PinotQueryGeneratorContext visitAggregation(AggregationNode node, PinotQu
                     case GROUP_BY: {
                         GroupByColumnNode groupByColumn = (GroupByColumnNode) expression;
                         VariableReferenceExpression groupByInputColumn = getVariableReference(groupByColumn.getInputColumn());
-                        checkState(groupByInputColumn.getType() instanceof FixedWidthType || groupByInputColumn.getType() instanceof VarcharType);
+                        checkState(groupByInputColumn.getType() instanceof FixedWidthType || groupByInputColumn.getType() instanceof VarcharType || groupByInputColumn.getType() instanceof JsonType);
                         variablesInAggregation.add(groupByInputColumn);
                         break;
                     }