Calcite UDF 实战之 ShardingSphere 联邦查询适配 MySQL BIT_COUNT

// bitwise
defineMethod(BITCOUNT, BuiltInMethod.BITCOUNT.method, NullPolicy.STRICT);
defineMethod(BIT_COUNT_BIG_QUERY, BuiltInMethod.BITCOUNT.method, NullPolicy.STRICT);
defineMethod(BIT_COUNT_MYSQL, BuiltInMethod.BITCOUNT.method, NullPolicy.STRICT);

@LibraryOperator(libraries = {MYSQL})
public static final SqlFunction BIT_COUNT_MYSQL =
    new SqlFunction(
        "BIT_COUNT",
        SqlKind.OTHER_FUNCTION,
        ReturnTypes.BIGINT_NULLABLE,
        null,
        OperandTypes.NUMERIC.or(OperandTypes.BINARY),
        SqlFunctionCategory.NUMERIC);

/**
 * Helper function for implementing <code>BITCOUNT</code>. Counts the number
 * of bits set in an integer value.
 */
public static long bitCount(long b) {
    return Long.bitCount(b);
}

private static final BigDecimal BITCOUNT_MAX = new BigDecimal(2).pow(64).subtract(new BigDecimal(1));
private static final BigDecimal BITCOUNT_MIN = new BigDecimal(2).pow(63).negate();

/**
 * Helper function for implementing <code>BITCOUNT</code>. Counts the number
 * of bits set in the integer portion of a decimal value.
 */
public static long bitCount(BigDecimal b) {
    final int comparison = b.compareTo(BITCOUNT_MAX);
    if (comparison < 0) {
        if (b.compareTo(BITCOUNT_MIN) <= 0) {
            return 1;
        } else {
            return bitCount(b.setScale(0, RoundingMode.DOWN).longValue());
        }
    } else if (comparison == 0) {
        return 64;
    } else {
        return 63;
    }
}

/**
 * Helper function for implementing <code>BITCOUNT</code>. Counts the number
 * of bits set in a ByteString value.
 */
public static long bitCount(ByteString b) {
    long bitsSet = 0;
    for (int i = 0; i < b.length(); i++) {
        bitsSet += Integer.bitCount(0xff & b.byteAt(i));
    }
    return bitsSet;
}

@LibraryOperator(libraries = {MYSQL})
public static final SqlFunction BIT_COUNT_MYSQL =
    new SqlFunction(
        "BIT_COUNT",
        SqlKind.OTHER_FUNCTION,
        ReturnTypes.BIGINT_NULLABLE,
        null,
        OperandTypes.or(OperandTypes.NUMERIC,
            OperandTypes.BINARY,
            OperandTypes.BOOLEAN,
            OperandTypes.CHARACTER,
            OperandTypes.DATETIME,
            OperandTypes.DATE,
            OperandTypes.TIME,
            OperandTypes.TIMESTAMP),
        SqlFunctionCategory.NUMERIC);

// 定义函数实现
defineMethod(BIT_COUNT_MYSQL, BuiltInMethod.BIT_COUNT_MYSQL.method, NullPolicy.STRICT);
// 声明 MySQL BIT_COUNT 函数实现方法
BIT_COUNT_MYSQL(SqlFunctions.class, "bitCountMySQL", Object.class),

/**
 * Helper function for implementing <code>BIT_COUNT_MYSQL</code>. Counts the number
 * of bits set in an Object value.
 */
public static long bitCountMySQL(Object b) {
    // Boolean 转换为 1 或 0 计算
    if (b instanceof Boolean) {
        return Long.bitCount((Boolean) b ? 1L : 0L);
    }
    // 当出现非数字字符时，捕获异常返回 0
    if (b instanceof String) {
        try {
            return bitCount(new BigDecimal((String) b));
        } catch (Exception ignore) {
            return 0;
        }
    }
    if (b instanceof Number) {
        return bitCount(new BigDecimal(b.toString()));
    }
    if (b instanceof ByteString) {
        return bitCount((ByteString) b);
    }
    // Date 去除非数字字符后计算
    if (b instanceof java.sql.Date) {
        return bitCountMySQL(new SimpleDateFormat("yyyyMMdd").format(((java.sql.Date) b)));
    }
    // Time 去除非数字字符后计算
    if (b instanceof Time) {
        return bitCountMySQL(new SimpleDateFormat("HHmmss").format(((Time) b)));
    }
    // Timestamp 去除非数字字符后计算
    if (b instanceof Timestamp) {
        return bitCountMySQL(new SimpleDateFormat("yyyyMMddHHmmss").format(((Timestamp) b)));
    }
    return 0;
}

// 执行 SELECT BIT_COUNT(joinedate) FROM EMPS_DATE_TIME LIMIT 1 生成代码
public org.apache.calcite.linq4j.Enumerable bind(final org.apache.calcite.DataContext root) {
  final org.apache.calcite.linq4j.Enumerable _inputEnumerable = org.apache.calcite.linq4j.Linq4j.asEnumerable(new Object[] {
    new Object[] {
      100,
      "Fred",
      10,
      null,
      null,
      40,
      Integer.valueOf(25),
      Boolean.valueOf(true),
      false,
      // DATE '1996-08-03' 转换为 9711
      9711},
    new Object[] {
      110,
      "Eric",
      20,
      "M",
      "San Francisco",
      3,
      Integer.valueOf(80),
      null,
      false,
      // DATE '2001-01-01' 转换为 11323
      11323},
    new Object[] {
      110,
      "John",
      40,
      "M",
      "Vancouver",
      2,
      null,
      Boolean.valueOf(false),
      true,
      // DATE '2002-05-03' 转换为 11810
      11810},
    new Object[] {
      120,
      "Wilma",
      20,
      "F",
      null,
      1,
      Integer.valueOf(5),
      null,
      true,
      // DATE '2005-09-07' 转换为 13033
      13033},
    new Object[] {
      130,
      "Alice",
      40,
      "F",
      "Vancouver",
      2,
      null,
      Boolean.valueOf(false),
      true,
      // DATE '2007-01-01' 转换为 13514
      13514}}).take(1);
  return new org.apache.calcite.linq4j.AbstractEnumerable(){
      public org.apache.calcite.linq4j.Enumerator enumerator() {
        return new org.apache.calcite.linq4j.Enumerator(){
            public final org.apache.calcite.linq4j.Enumerator inputEnumerator = _inputEnumerable.enumerator();
            public void reset() {
              inputEnumerator.reset();
            }

            public boolean moveNext() {
              return inputEnumerator.moveNext();
            }

            public void close() {
              inputEnumerator.close();
            }

            public Object current() {
              // 先调用 SqlFunctions.toInt，将 long 转为 int，再调用 SqlFunctions.bitCountMySQL
              return org.apache.calcite.runtime.SqlFunctions.bitCountMySQL(org.apache.calcite.runtime.SqlFunctions.toInt(((Object[]) inputEnumerator.current())[9]));
            }
          };
      }
    };
}

public Class getElementType() {
  return long.class;
}

/**
 * Visit {@code RexCall}. For most {@code SqlOperator}s, we can get the implementor
 * from {@code RexImpTable}. Several operators (e.g., CaseWhen) with special semantics
 * need to be implemented separately.
 */
@Override
public Result visitCall(RexCall call) {
    if (rexResultMap.containsKey(call)) {
        return rexResultMap.get(call);
    }
    ...
    // 从 RexImpTable.INSTANCE 中获取函数实现器 Implementor
    final RexImpTable.RexCallImplementor implementor = RexImpTable.INSTANCE.get(operator);
    if (implementor == null) {
        throw new RuntimeException("cannot translate call " + call);
    }
    // 获取函数操作数
    final List<RexNode> operandList = call.getOperands();
    // 转换为内部类型
    final List<@Nullable Type> storageTypes = EnumUtils.internalTypes(operandList);
    final List<Result> operandResults = new ArrayList<>();
    // 遍历操作数，实现函数逻辑
    for (int i = 0; i < operandList.size(); i++) {
        final Result operandResult = implementCallOperand(operandList.get(i), storageTypes.get(i), this);
        operandResults.add(operandResult);
    }
    callOperandResultMap.put(call, operandResults);
    final Result result = implementor.implement(this, call, operandResults);
    rexResultMap.put(call, result);
    return result;
}

// EnumUtils#toInternal 方法
static @Nullable Type toInternal(RelDataType type, boolean forceNotNull) {
    switch (type.getSqlTypeName()) {
        case DATE:
        case TIME:
            return type.isNullable() && !forceNotNull ? Integer.class : int.class;
        case TIMESTAMP:
            return type.isNullable() && !forceNotNull ? Long.class : long.class;
        default:
            return null; // we don't care; use the default storage type
    }
}

/**
 * Implementor for MYSQL {@code BIT_COUNT} function.
 */
private static class BitCountMySQLImplementor extends AbstractRexCallImplementor {
    BitCountMySQLImplementor() {
        super("bitCount", NullPolicy.STRICT, false);
    }
    
    @Override
    Expression implementSafe(final RexToLixTranslator translator, final RexCall call, final List<Expression> argValueList) {
        Expression expr = argValueList.get(0);
        RelDataType relDataType = call.getOperands().get(0).getType();
        if (SqlTypeUtil.isNull(relDataType)) {
            return argValueList.get(0);
        }
        // In MySQL, BIT_COUNT(TIMESTAMP '1996-08-03 16:22:34') is converted to
        // BIT_COUNT('19960803162234') for calculation, so the internal int value
        // needs to be converted to DATE/TIME and TIMESTAMP.
        SqlTypeName type = relDataType.getSqlTypeName();
        switch (type) {
            case VARBINARY:
            case BINARY:
                return Expressions.call(SqlFunctions.class, "bitCount", expr);
            case DATE:
                expr = Expressions.call(BuiltInMethod.INTERNAL_TO_DATE.method, expr);
                break;
            case TIME:
                expr = Expressions.call(BuiltInMethod.INTERNAL_TO_TIME.method, expr);
                break;
            case TIMESTAMP:
                expr = Expressions.call(BuiltInMethod.INTERNAL_TO_TIMESTAMP.method, expr);
                break;
            default:
                break;
        }
        return Expressions.call(SqlFunctions.class, "bitCountMySQL", Expressions.box(expr));
    }
}

define(BIT_COUNT_MYSQL, new BitCountMySQLImplementor());