ma.vi.esql.parser.macro.Bin Maven / Gradle / Ivy

/*
 * Copyright (c) 2020 Vikash Madhow
 */

package ma.vi.esql.parser.macro;

import ma.vi.esql.function.Function;
import ma.vi.esql.function.FunctionParameter;
import ma.vi.esql.parser.Context;
import ma.vi.esql.parser.Esql;
import ma.vi.esql.parser.Macro;
import ma.vi.esql.parser.TranslationException;
import ma.vi.esql.parser.expression.*;
import ma.vi.esql.type.Types;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Iterator;
import java.util.List;

import static java.lang.String.valueOf;
import static org.apache.commons.lang3.StringUtils.leftPad;

/**
 * Given a value and a variadic array of intervals, returns the range
 * where the value fits. E.g.
 * 
 * bin(15, 'x', 1,5,12,20,35,67) returns '12 <= x < 19'
 * 

 * bin is macro which is expanded to a case statement: E.g.
 * 
 *   bin(age, 'age', 1,2,5) expands to:
 *
 *      '01. age < 1'         if age < 1         else
 *      '02. 1 <= age < 2' if 1 <= age < 2 else
 *      '03. 2 <= age < 5' if 2 <= age < 5 else '04. age >= 5'
 *
 *   which is then translated to a database-specific CASE expression.
 * 
 *
 * The interval values can be expressions but they must be in correct order,
 * i.e., from smallest to largest. Otherwise the case expansion will be wrong.
 *
 * @author Vikash Madhow ([email protected])
 */
public class Bin extends Function implements Macro {
  public Bin() {
    super("bin", Types.TextType,
          Arrays.asList(new FunctionParameter("val", Types.TextType),
                        new FunctionParameter("name", Types.TextType),
                        new FunctionParameter("intervals", Types.TextType)));
  }

  @Override
  public boolean expand(String name, Esql esql) {
    FunctionCall call = (FunctionCall)esql;
    Context ctx = call.context;
    List> args = call.arguments();
    if (args.size() < 3) {
      throw new TranslationException("bin requires at least 3 arguments: the value to bin, a human-readable "
                                   + "name of the value to bin, and at least 1 value defining the intervals of "
                                   + "the bin range. E.g: bin(x, 'age', 1, 5, 10) will produce bins: "
                                   + "age < 1, 1 <= age < 5, 5 <= age < 10 and age >= 10" );
    }

    int order = 2;
    List> cases = new ArrayList<>();
    Iterator> i = args.iterator();
    Expression binVar = i.next();
    Expression varName = i.next();
    Expression lower = null;
    Expression upper = i.next();
    while (upper != null) {
      if (lower == null) {
        cases.add(new Concatenation(ctx,
                                    Arrays.asList(
                                        new StringLiteral(ctx, "01. "),
                                        (Expression)varName,
                                        new StringLiteral(ctx, " < "),
                                        new Cast(ctx, upper, Types.StringType))));
        cases.add(new LessThan(ctx, (Expression)binVar, (Expression)upper));
      }
      lower = upper;
      upper = i.hasNext() ? i.next() : null;
      if (upper == null) {
        cases.add(new Concatenation(ctx,
                                    Arrays.asList(
                                        new StringLiteral(ctx, leftPad(valueOf(order), 2, '0') + ". "),
                                        (Expression)varName,
                                        new StringLiteral(ctx, " >= "),
                                        new Cast(ctx, lower, Types.StringType))));
      } else {
        cases.add(new Concatenation(ctx,
                                    Arrays.asList(
                                        new StringLiteral(ctx, leftPad(valueOf(order), 2, '0') + ". "),
                                        new Cast(ctx, lower, Types.StringType),
                                        new StringLiteral(ctx," <= "),
                                        (Expression)varName,
                                        new StringLiteral(ctx," < "),
                                        new Cast(ctx, upper, Types.StringType))));
        cases.add(new Range(ctx,
                            (Expression)lower,
                            "<=",
                            (Expression)binVar,
                            "<",
                            (Expression)upper));
        order += 1;
      }
    }
    call.parent.replaceWith(name, new Case(ctx, cases));
    return true;
  }
}