Examples with ColumnRef com.pingcap.tikv.expression.ColumnRef used on opensource projects

Example 1 with ColumnRef

use of com.pingcap.tikv.expression.ColumnRef in project tispark by pingcap.

the class IndexRangeSetBuilder method visit.

@Override
protected RangeSet<TypedKey> visit(StringRegExpression node, Void context) {
    ColumnRef columnRef = node.getColumnRef();
    // In order to match a prefix index, we have to cut the literal by prefix length.
    // e.g., for table t:
    // CREATE TABLE `t` {
    // `c1` VARCHAR(10) DEFAULT NULL,
    // KEY `prefix_index` (`c`(2))
    // }
    // when the predicate is `c1` LIKE 'abc%', the index range should be ['ab', 'ab'].
    // when the predicate is `c1` LIKE 'a%', the index range should be ['a', 'b').
    // for varchar, `c1`(2) will take first two characters(bytes) as prefix index.
    // TODO: Note that TiDB only supports UTF-8, we need to check if prefix index behave differently
    // under other encoding methods
    int prefixLen = lengths.getOrDefault(columnRef, DataType.UNSPECIFIED_LEN);
    TypedKey literal = node.getTypedLiteral(prefixLen);
    RangeSet<TypedKey> ranges = TreeRangeSet.create();
    switch(node.getRegType()) {
        case STARTS_WITH:
            ranges.add(Range.atLeast(literal).intersection(Range.lessThan(literal.next())));
            break;
        default:
            throwOnError(node);
    }
    return ranges;
}

Example 2 with ColumnRef

use of com.pingcap.tikv.expression.ColumnRef in project tispark by pingcap.

the class TiDAGRequest method buildScan.

/**
 * Unify indexScan and tableScan building logic since they are very much alike. DAGRequest for
 * IndexScan should also contain filters and aggregation, so we can reuse this part of logic.
 *
 * <p>DAGRequest is made up of a chain of executors with strict orders: TableScan/IndexScan >
 * Selection > Aggregation > TopN/Limit a DAGRequest must contain one and only one TableScan or
 * IndexScan.
 *
 * @param buildIndexScan whether the dagRequest to build should be an {@link
 *     com.pingcap.tidb.tipb.IndexScan}
 * @return final DAGRequest built
 */
private DAGRequest.Builder buildScan(boolean buildIndexScan, List<Integer> outputOffsets) {
    long id = getPhysicalId();
    checkNotNull(startTs, "startTs is null");
    checkArgument(startTs.getVersion() != 0, "timestamp is 0");
    clearPushDownInfo();
    DAGRequest.Builder dagRequestBuilder = DAGRequest.newBuilder();
    Executor.Builder executorBuilder = Executor.newBuilder();
    IndexScan.Builder indexScanBuilder = IndexScan.newBuilder();
    TableScan.Builder tblScanBuilder = TableScan.newBuilder();
    // find a column's offset in fields
    Map<String, Integer> colOffsetInFieldMap = new HashMap<>();
    // find a column's position in index
    Map<String, Integer> colPosInIndexMap = new HashMap<>();
    if (buildIndexScan) {
        // IndexScan
        if (indexInfo == null) {
            throw new TiClientInternalException("Index is empty for index scan");
        }
        List<TiColumnInfo> columnInfoList = tableInfo.getColumns();
        boolean hasPk = false;
        // We extract index column info
        List<Integer> indexColOffsets = indexInfo.getIndexColumns().stream().map(TiIndexColumn::getOffset).collect(Collectors.toList());
        int idxPos = 0;
        // for index scan builder, columns are added by its order in index
        for (Integer idx : indexColOffsets) {
            TiColumnInfo tiColumnInfo = columnInfoList.get(idx);
            ColumnInfo columnInfo = tiColumnInfo.toProto(tableInfo);
            colPosInIndexMap.put(tiColumnInfo.getName(), idxPos++);
            ColumnInfo.Builder colBuilder = ColumnInfo.newBuilder(columnInfo);
            if (columnInfo.getColumnId() == -1) {
                hasPk = true;
                colBuilder.setPkHandle(true);
            }
            indexScanBuilder.addColumns(colBuilder);
        }
        int colCount = indexScanBuilder.getColumnsCount();
        if (isDoubleRead()) {
            // TODO: we may merge indexDoubleRead and coveringIndexRead logic
            for (ColumnRef col : getFields()) {
                Integer pos = colPosInIndexMap.get(col.getName());
                if (pos != null) {
                    TiColumnInfo columnInfo = columnInfoList.get(indexColOffsets.get(pos));
                    if (col.matchName(columnInfo.getName())) {
                        colOffsetInFieldMap.put(col.getName(), pos);
                    }
                // TODO: primary key may also be considered if pkIsHandle
                }
            }
            // double read case
            if (!hasPk) {
                // add handle column
                if (!tableInfo.isCommonHandle()) {
                    indexScanBuilder.addColumns(handleColumn);
                    ++colCount;
                } else {
                    for (TiIndexColumn col : tableInfo.getPrimaryKey().getIndexColumns()) {
                        indexScanBuilder.addColumns(tableInfo.getColumn(col.getName()).toProto(tableInfo));
                        ++colCount;
                    }
                }
                addRequiredIndexDataType();
            }
            if (colCount == 0) {
                throw new DAGRequestException("Incorrect index scan with zero column count");
            }
            if (!tableInfo.isCommonHandle()) {
                outputOffsets.add(colCount - 1);
            } else {
                int idxColSize = tableInfo.getPrimaryKey().getIndexColumns().size();
                for (int i = idxColSize; i >= 1; i--) {
                    outputOffsets.add(colCount - i);
                }
            }
        } else {
            boolean pkIsNeeded = false;
            // offset for dagRequest should be in accordance with fields
            for (ColumnRef col : getFields()) {
                Integer pos = colPosInIndexMap.get(col.getName());
                if (pos != null) {
                    TiColumnInfo columnInfo = columnInfoList.get(indexColOffsets.get(pos));
                    if (col.matchName(columnInfo.getName())) {
                        outputOffsets.add(pos);
                        colOffsetInFieldMap.put(col.getName(), pos);
                    }
                } else // logically it must be the pk column. Extra check here.
                if (tableInfo.getColumn(col.getName()).isPrimaryKey()) {
                    pkIsNeeded = true;
                    // offset should be processed for each primary key encountered
                    outputOffsets.add(colCount);
                    // for index scan, column offset must be in the order of index->handle
                    colOffsetInFieldMap.put(col.getName(), indexColOffsets.size());
                } else {
                    throw new DAGRequestException("columns other than primary key and index key exist in fields while index single read: " + col.getName());
                }
            }
            // pk is not included in index but still needed
            if (pkIsNeeded) {
                if (!tableInfo.isCommonHandle()) {
                    indexScanBuilder.addColumns(handleColumn);
                }
            }
        }
        executorBuilder.setTp(ExecType.TypeIndexScan);
        indexScanBuilder.setTableId(id).setIndexId(indexInfo.getId());
        if (tableInfo.isCommonHandle()) {
            for (TiIndexColumn col : tableInfo.getPrimaryKey().getIndexColumns()) {
                indexScanBuilder.addPrimaryColumnIds(tableInfo.getColumn(col.getName()).getId());
            }
        }
        dagRequestBuilder.addExecutors(executorBuilder.setIdxScan(indexScanBuilder).build());
    } else {
        // TableScan
        executorBuilder.setTp(ExecType.TypeTableScan);
        tblScanBuilder.setTableId(id);
        if (tableInfo.isCommonHandle()) {
            for (TiIndexColumn col : tableInfo.getPrimaryKey().getIndexColumns()) {
                tblScanBuilder.addPrimaryColumnIds(tableInfo.getColumn(col.getName()).getId());
            }
        }
        // Step1. Add columns to first executor
        int lastOffset = 0;
        for (ColumnRef col : getFields()) {
            // can't allow duplicated col added into executor.
            if (!colOffsetInFieldMap.containsKey(col.getName())) {
                tblScanBuilder.addColumns(tableInfo.getColumn(col.getName()).toProto(tableInfo));
                colOffsetInFieldMap.put(col.getName(), lastOffset);
                lastOffset++;
            }
            // column offset should be in accordance with fields
            outputOffsets.add(colOffsetInFieldMap.get(col.getName()));
        }
        dagRequestBuilder.addExecutors(executorBuilder.setTblScan(tblScanBuilder));
    }
    boolean isIndexDoubleScan = buildIndexScan && isDoubleRead();
    // Should build these executors when performing CoveringIndexScan/TableScan
    // clear executorBuilder
    executorBuilder.clear();
    // Step2. Add others
    // DO NOT EDIT EXPRESSION CONSTRUCTION ORDER
    // Or make sure the construction order is below:
    // TableScan/IndexScan > Selection > Aggregation > TopN/Limit
    Expression whereExpr = mergeCNFExpressions(getFilters());
    if (whereExpr != null) {
        if (!isIndexDoubleScan || isExpressionCoveredByIndex(whereExpr)) {
            executorBuilder.setTp(ExecType.TypeSelection);
            dagRequestBuilder.addExecutors(executorBuilder.setSelection(Selection.newBuilder().addConditions(ProtoConverter.toProto(whereExpr, colOffsetInFieldMap))));
            executorBuilder.clear();
            addPushDownFilters();
        } else {
            return dagRequestBuilder;
        }
    }
    if (!getGroupByItems().isEmpty() || !getAggregates().isEmpty()) {
        // only allow table scan or covering index scan push down groupby and agg
        if (!isIndexDoubleScan || (isGroupByCoveredByIndex() && isAggregateCoveredByIndex())) {
            pushDownAggAndGroupBy(dagRequestBuilder, executorBuilder, outputOffsets, colOffsetInFieldMap);
        } else {
            return dagRequestBuilder;
        }
    }
    if (!getOrderByItems().isEmpty()) {
        if (!isIndexDoubleScan || isOrderByCoveredByIndex()) {
            // only allow table scan or covering index scan push down orderby
            pushDownOrderBy(dagRequestBuilder, executorBuilder, colOffsetInFieldMap);
        }
    } else if (getLimit() != 0) {
        if (!isIndexDoubleScan) {
            pushDownLimit(dagRequestBuilder, executorBuilder);
        }
    }
    return dagRequestBuilder;
}

Example 3 with ColumnRef

use of com.pingcap.tikv.expression.ColumnRef in project tispark by pingcap.

the class TiDAGRequest method addRequiredIndexDataType.

/**
 * Required index columns for double read
 */
private void addRequiredIndexDataType() {
    if (!tableInfo.isCommonHandle()) {
        indexDataTypes.add(requireNonNull(IntegerType.BIGINT, "dataType is null"));
    } else {
        for (TiIndexColumn col : tableInfo.getPrimaryKey().getIndexColumns()) {
            String c = col.getName();
            ColumnRef cr = ColumnRef.create(c, tableInfo.getColumn(c));
            indexDataTypes.add(cr.getDataType());
        }
    }
}

Example 4 with ColumnRef

use of com.pingcap.tikv.expression.ColumnRef in project tispark by pingcap.

the class PredicateUtilsTest method expressionToIndexRangesTest.

@Test
public void expressionToIndexRangesTest() {
    TiTableInfo table = createTable();
    ColumnRef col1 = ColumnRef.create("c1", table);
    ColumnRef col4 = ColumnRef.create("c4", table);
    ColumnRef col5 = ColumnRef.create("c5", table);
    Constant c1 = Constant.create(1, IntegerType.INT);
    Constant c2 = Constant.create(2, IntegerType.INT);
    Constant c3 = Constant.create(3, IntegerType.INT);
    Constant c4 = Constant.create(4, IntegerType.INT);
    TypedKey key1 = TypedKey.toTypedKey(1, IntegerType.INT);
    TypedKey key2 = TypedKey.toTypedKey(2, IntegerType.INT);
    TypedKey key3 = TypedKey.toTypedKey(3, IntegerType.INT);
    TypedKey key4 = TypedKey.toTypedKey(4, IntegerType.INT);
    Expression predicate1 = or(or(equal(c1, col1), equal(col1, c2)), equal(col1, c1));
    Expression predicate2 = or(equal(c3, col4), equal(c4, col4));
    Expression rangePredicate = notEqual(col5, c1);
    List<IndexRange> indexRanges = PredicateUtils.expressionToIndexRanges(ImmutableList.of(predicate1, predicate2), Optional.of(rangePredicate), table, null);
    assertEquals(8, indexRanges.size());
    Key indexKey1 = CompoundKey.concat(key1, key3);
    Key indexKey2 = CompoundKey.concat(key1, key4);
    Key indexKey3 = CompoundKey.concat(key2, key3);
    Key indexKey4 = CompoundKey.concat(key2, key4);
    Range<TypedKey> baselineRange1 = Range.lessThan(key1);
    Range<TypedKey> baselineRange2 = Range.greaterThan(key1);
    Set<Key> baselineKeys = ImmutableSet.of(indexKey1, indexKey2, indexKey3, indexKey4);
    Set<Range<TypedKey>> baselineRanges = ImmutableSet.of(baselineRange1, baselineRange2);
    for (IndexRange range : indexRanges) {
        assertTrue(baselineKeys.contains(range.getAccessKey()));
        assertTrue(baselineRanges.contains(range.getRange()));
    }
}

Example 5 with ColumnRef

use of com.pingcap.tikv.expression.ColumnRef in project tispark by pingcap.

the class TiDAGRequestTest method testTopNCouldPushDownLimit0.

@Test
public void testTopNCouldPushDownLimit0() {
    TiTableInfo table = createTable();
    TiDAGRequest dagRequest = new TiDAGRequest(TiDAGRequest.PushDownType.NORMAL);
    ColumnRef col1 = ColumnRef.create("c1", table);
    dagRequest.addOrderByItem(ByItem.create(col1, false));
    dagRequest.addRequiredColumn(col1);
    dagRequest.setLimit(0);
    dagRequest.setTableInfo(table);
    dagRequest.setStartTs(new TiTimestamp(0, 1));
    dagRequest.buildTableScan();
}