Skip to content

Commit 6510d2b

Browse files
srikanthpadakantisrikanthpadakanti
committed
Address coderrabbit comments
Signed-off-by: Srikanth Padakanti <srikanth_padakanti@apple.com>
1 parent b2a9d9f commit 6510d2b

1 file changed

Lines changed: 207 additions & 40 deletions

File tree

‎core/src/main/java/org/opensearch/sql/calcite/CalciteRelNodeVisitor.java‎

Lines changed: 207 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -3118,6 +3118,139 @@ public RelNode visitExpand(Expand expand, CalcitePlanContext context) {
31183118
return context.relBuilder.peek();
31193119
}
31203120

3121+
// /**
3122+
// * mvcombine command visitor to collapse rows that are identical on all fields except the
3123+
// target
3124+
// * field, and combine the target field values into a multivalue (array) field.
3125+
// *
3126+
// *
3127+
// * @param node mvcombine command to be visited
3128+
// * @param context CalcitePlanContext containing the RelBuilder, RexBuilder, and resolution
3129+
// context
3130+
// * @return RelNode representing records collapsed by all non-target fields with the target
3131+
// field
3132+
// * combined into a multivalue array (or a delimited string when {@code nomv=true})
3133+
// */
3134+
// @Override
3135+
// public RelNode visitMvCombine(MvCombine node, CalcitePlanContext context) {
3136+
// // 1) Lower child first
3137+
// visitChildren(node, context);
3138+
//
3139+
// final RelBuilder relBuilder = context.relBuilder;
3140+
// final RexBuilder rexBuilder = context.rexBuilder;
3141+
//
3142+
// final RelNode input = relBuilder.peek();
3143+
// final List<String> inputFieldNames = input.getRowType().getFieldNames();
3144+
//
3145+
// // 2) Accept delim (default is single space).
3146+
// // NOTE: delim only affects single-value rendering when nomv=true.
3147+
// final String delim = node.getDelim();
3148+
//
3149+
// // 3) Resolve target field to an input ref (must be a direct field reference)
3150+
// final Field targetField = node.getField();
3151+
// final RexNode targetRex = rexVisitor.analyze(targetField, context);
3152+
// if (!(targetRex instanceof RexInputRef)) {
3153+
// throw new SemanticCheckException(
3154+
// "mvcombine target must be a direct field reference, but got: " + targetField);
3155+
// }
3156+
// final int targetIndex = ((RexInputRef) targetRex).getIndex();
3157+
// final String targetName = inputFieldNames.get(targetIndex);
3158+
//
3159+
// // 4) Group key = all fields except target
3160+
// final List<RexNode> groupExprs = new ArrayList<>();
3161+
// for (int i = 0; i < inputFieldNames.size(); i++) {
3162+
// if (i == targetIndex) continue;
3163+
// groupExprs.add(relBuilder.field(i));
3164+
// }
3165+
//
3166+
// // 5) Aggregate target values: COLLECT => MULTISET
3167+
// final RelBuilder.AggCall aggCall =
3168+
// relBuilder
3169+
// .aggregateCall(SqlStdOperatorTable.COLLECT, relBuilder.field(targetIndex))
3170+
// .as(targetName);
3171+
//
3172+
// relBuilder.aggregate(relBuilder.groupKey(groupExprs), aggCall);
3173+
//
3174+
// // 6) Restore original output column order AND cast MULTISET -> ARRAY using RexBuilder
3175+
// // After aggregate => [group fields..., targetAgg(multiset)]
3176+
// final int aggPos = groupExprs.size();
3177+
//
3178+
// final RelDataType elemType = input.getRowType().getFieldList().get(targetIndex).getType();
3179+
// final RelDataType arrayType = relBuilder.getTypeFactory().createArrayType(elemType, -1);
3180+
//
3181+
// final List<RexNode> projections = new ArrayList<>(inputFieldNames.size());
3182+
// final List<String> projNames = new ArrayList<>(inputFieldNames.size());
3183+
//
3184+
// int groupPos = 0;
3185+
// for (int i = 0; i < inputFieldNames.size(); i++) {
3186+
// projNames.add(inputFieldNames.get(i));
3187+
// if (i == targetIndex) {
3188+
// final RexNode multisetRef = relBuilder.field(aggPos); // MULTISET
3189+
// projections.add(rexBuilder.makeCast(arrayType, multisetRef)); // ARRAY
3190+
// } else {
3191+
// projections.add(relBuilder.field(groupPos));
3192+
// groupPos++;
3193+
// }
3194+
// }
3195+
// relBuilder.project(projections, projNames, /* force= */ true);
3196+
//
3197+
// // 7) nomv=true converts multivalue output to a single delimited string.
3198+
// // arrayToString in this engine supports only String/ByteString, so cast elements to STRING
3199+
// // first.
3200+
// if (node.isNomv()) {
3201+
// final List<RexNode> nomvProjections = new ArrayList<>(inputFieldNames.size());
3202+
// final List<String> nomvNames = new ArrayList<>(inputFieldNames.size());
3203+
//
3204+
// // Build ARRAY<VARCHAR> type once
3205+
// final RelDataType stringType =
3206+
// relBuilder
3207+
// .getTypeFactory()
3208+
// .createSqlType(org.apache.calcite.sql.type.SqlTypeName.VARCHAR);
3209+
// final RelDataType stringArrayType =
3210+
// relBuilder.getTypeFactory().createArrayType(stringType, -1);
3211+
//
3212+
// for (int i = 0; i < inputFieldNames.size(); i++) {
3213+
// final String name = inputFieldNames.get(i);
3214+
// nomvNames.add(name);
3215+
//
3216+
// if (i == targetIndex) {
3217+
// // At this point relBuilder.field(i) is ARRAY<elemType>. Cast to ARRAY<VARCHAR> so
3218+
// // ARRAY_TO_STRING works.
3219+
// final RexNode stringArray = rexBuilder.makeCast(stringArrayType, relBuilder.field(i));
3220+
//
3221+
// nomvProjections.add(
3222+
// relBuilder.call(
3223+
// org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_TO_STRING,
3224+
// stringArray,
3225+
// relBuilder.literal(delim)));
3226+
// } else {
3227+
// nomvProjections.add(relBuilder.field(i));
3228+
// }
3229+
// }
3230+
//
3231+
// relBuilder.project(nomvProjections, nomvNames, /* force= */ true);
3232+
// }
3233+
//
3234+
// return relBuilder.peek();
3235+
// }
3236+
3237+
/**
3238+
* mvcombine command visitor to collapse rows that are identical on all fields except the target
3239+
* field, and combine the target field values into a multivalue (array) field.
3240+
*
3241+
* <p>Implementation notes:
3242+
*
3243+
* <p>Groups by all input fields except the target field. Aggregates target using {@code COLLECT}
3244+
* (MULTISET) and casts MULTISET -> ARRAY. Preserves original output column order. If {@code
3245+
* nomv=true}, renders target array as a scalar string via {@code ARRAY_TO_STRING} using {@code
3246+
* delim}. {@code delim} only affects output when {@code nomv=true}.
3247+
*
3248+
* @param node mvcombine command to be visited
3249+
* @param context CalcitePlanContext containing the RelBuilder, RexBuilder, and resolution context
3250+
* @return RelNode representing collapsed records with the target combined into a multivalue array
3251+
* (or a delimited string when {@code nomv=true})
3252+
* @throws SemanticCheckException if the mvcombine target is not a direct field reference
3253+
*/
31213254
@Override
31223255
public RelNode visitMvCombine(MvCombine node, CalcitePlanContext context) {
31233256
// 1) Lower child first
@@ -3135,30 +3268,65 @@ public RelNode visitMvCombine(MvCombine node, CalcitePlanContext context) {
31353268

31363269
// 3) Resolve target field to an input ref (must be a direct field reference)
31373270
final Field targetField = node.getField();
3271+
final int targetIndex = resolveTargetIndex(targetField, context);
3272+
final String targetName = inputFieldNames.get(targetIndex);
3273+
3274+
// 4) Group key = all fields except target
3275+
final List<RexNode> groupExprs =
3276+
buildGroupExpressionsExcludingTarget(targetIndex, inputFieldNames, relBuilder);
3277+
3278+
// 5) Aggregate target values: COLLECT => MULTISET
3279+
performCollectAggregation(relBuilder, targetIndex, targetName, groupExprs);
3280+
3281+
// 6) Restore original output column order AND cast MULTISET -> ARRAY using RexBuilder
3282+
restoreColumnOrderWithArrayCast(
3283+
relBuilder, rexBuilder, input, inputFieldNames, targetIndex, groupExprs);
3284+
3285+
// 7) nomv=true converts multivalue output to a single delimited string.
3286+
if (node.isNomv()) {
3287+
applyNomvConversion(relBuilder, rexBuilder, inputFieldNames, targetIndex, delim);
3288+
}
3289+
3290+
return relBuilder.peek();
3291+
}
3292+
3293+
private int resolveTargetIndex(Field targetField, CalcitePlanContext context) {
31383294
final RexNode targetRex = rexVisitor.analyze(targetField, context);
31393295
if (!(targetRex instanceof RexInputRef)) {
31403296
throw new SemanticCheckException(
31413297
"mvcombine target must be a direct field reference, but got: " + targetField);
31423298
}
3143-
final int targetIndex = ((RexInputRef) targetRex).getIndex();
3144-
final String targetName = inputFieldNames.get(targetIndex);
3299+
return ((RexInputRef) targetRex).getIndex();
3300+
}
31453301

3146-
// 4) Group key = all fields except target
3302+
private List<RexNode> buildGroupExpressionsExcludingTarget(
3303+
int targetIndex, List<String> inputFieldNames, RelBuilder relBuilder) {
31473304
final List<RexNode> groupExprs = new ArrayList<>();
31483305
for (int i = 0; i < inputFieldNames.size(); i++) {
31493306
if (i == targetIndex) continue;
31503307
groupExprs.add(relBuilder.field(i));
31513308
}
3309+
return groupExprs;
3310+
}
31523311

3153-
// 5) Aggregate target values: COLLECT => MULTISET
3312+
private void performCollectAggregation(
3313+
RelBuilder relBuilder, int targetIndex, String targetName, List<RexNode> groupExprs) {
31543314
final RelBuilder.AggCall aggCall =
31553315
relBuilder
31563316
.aggregateCall(SqlStdOperatorTable.COLLECT, relBuilder.field(targetIndex))
31573317
.as(targetName);
31583318

31593319
relBuilder.aggregate(relBuilder.groupKey(groupExprs), aggCall);
3320+
}
3321+
3322+
private void restoreColumnOrderWithArrayCast(
3323+
RelBuilder relBuilder,
3324+
RexBuilder rexBuilder,
3325+
RelNode input,
3326+
List<String> inputFieldNames,
3327+
int targetIndex,
3328+
List<RexNode> groupExprs) {
31603329

3161-
// 6) Restore original output column order AND cast MULTISET -> ARRAY using RexBuilder
31623330
// After aggregate => [group fields..., targetAgg(multiset)]
31633331
final int aggPos = groupExprs.size();
31643332

@@ -3179,46 +3347,45 @@ public RelNode visitMvCombine(MvCombine node, CalcitePlanContext context) {
31793347
groupPos++;
31803348
}
31813349
}
3350+
31823351
relBuilder.project(projections, projNames, /* force= */ true);
3352+
}
31833353

3184-
// 7) nomv=true converts multivalue output to a single delimited string.
3185-
// arrayToString in this engine supports only String/ByteString, so cast elements to STRING
3186-
// first.
3187-
if (node.isNomv()) {
3188-
final List<RexNode> nomvProjections = new ArrayList<>(inputFieldNames.size());
3189-
final List<String> nomvNames = new ArrayList<>(inputFieldNames.size());
3190-
3191-
// Build ARRAY<VARCHAR> type once
3192-
final RelDataType stringType =
3193-
relBuilder
3194-
.getTypeFactory()
3195-
.createSqlType(org.apache.calcite.sql.type.SqlTypeName.VARCHAR);
3196-
final RelDataType stringArrayType =
3197-
relBuilder.getTypeFactory().createArrayType(stringType, -1);
3198-
3199-
for (int i = 0; i < inputFieldNames.size(); i++) {
3200-
final String name = inputFieldNames.get(i);
3201-
nomvNames.add(name);
3202-
3203-
if (i == targetIndex) {
3204-
// At this point relBuilder.field(i) is ARRAY<elemType>. Cast to ARRAY<VARCHAR> so
3205-
// ARRAY_TO_STRING works.
3206-
final RexNode stringArray = rexBuilder.makeCast(stringArrayType, relBuilder.field(i));
3207-
3208-
nomvProjections.add(
3209-
relBuilder.call(
3210-
org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_TO_STRING,
3211-
stringArray,
3212-
relBuilder.literal(delim)));
3213-
} else {
3214-
nomvProjections.add(relBuilder.field(i));
3215-
}
3216-
}
3354+
private void applyNomvConversion(
3355+
RelBuilder relBuilder,
3356+
RexBuilder rexBuilder,
3357+
List<String> inputFieldNames,
3358+
int targetIndex,
3359+
String delim) {
32173360

3218-
relBuilder.project(nomvProjections, nomvNames, /* force= */ true);
3361+
final List<RexNode> nomvProjections = new ArrayList<>(inputFieldNames.size());
3362+
final List<String> nomvNames = new ArrayList<>(inputFieldNames.size());
3363+
3364+
// Build ARRAY<VARCHAR> type once
3365+
final RelDataType stringType =
3366+
relBuilder.getTypeFactory().createSqlType(org.apache.calcite.sql.type.SqlTypeName.VARCHAR);
3367+
final RelDataType stringArrayType = relBuilder.getTypeFactory().createArrayType(stringType, -1);
3368+
3369+
for (int i = 0; i < inputFieldNames.size(); i++) {
3370+
final String name = inputFieldNames.get(i);
3371+
nomvNames.add(name);
3372+
3373+
if (i == targetIndex) {
3374+
// At this point relBuilder.field(i) is ARRAY<elemType>. Cast to ARRAY<VARCHAR> so
3375+
// ARRAY_TO_STRING works.
3376+
final RexNode stringArray = rexBuilder.makeCast(stringArrayType, relBuilder.field(i));
3377+
3378+
nomvProjections.add(
3379+
relBuilder.call(
3380+
org.apache.calcite.sql.fun.SqlLibraryOperators.ARRAY_TO_STRING,
3381+
stringArray,
3382+
relBuilder.literal(delim)));
3383+
} else {
3384+
nomvProjections.add(relBuilder.field(i));
3385+
}
32193386
}
32203387

3221-
return relBuilder.peek();
3388+
relBuilder.project(nomvProjections, nomvNames, /* force= */ true);
32223389
}
32233390

32243391
@Override

0 commit comments

Comments
 (0)