io.prestosql.spi.connector.RecordCursor Java Examples

@Test
public void test()
{
    ArrayType arrayOfBigintType = new ArrayType(BIGINT);
    FieldSetFilteringRecordSet fieldSetFilteringRecordSet = new FieldSetFilteringRecordSet(
            createTestMetadataManager(),
            new InMemoryRecordSet(
                    ImmutableList.of(BIGINT, BIGINT, TIMESTAMP_WITH_TIME_ZONE, TIMESTAMP_WITH_TIME_ZONE, arrayOfBigintType, arrayOfBigintType),
                    ImmutableList.of(
                            ImmutableList.of(
                                    100L,
                                    100L,
                                    // test same time in different time zone to make sure equal check was done properly
                                    packDateTimeWithZone(100, getTimeZoneKeyForOffset(123)),
                                    packDateTimeWithZone(100, getTimeZoneKeyForOffset(234)),
                                    // test structural type
                                    arrayBlockOf(BIGINT, 12, 34, 56),
                                    arrayBlockOf(BIGINT, 12, 34, 56)))),
            ImmutableList.of(ImmutableSet.of(0, 1), ImmutableSet.of(2, 3), ImmutableSet.of(4, 5)));
    RecordCursor recordCursor = fieldSetFilteringRecordSet.cursor();
    assertTrue(recordCursor.advanceNextPosition());
}
 

public static PrestoThriftBlock fromLongBasedColumn(RecordSet recordSet, int columnIndex, int positions, BiFunction<boolean[], long[], PrestoThriftBlock> result)
{
    if (positions == 0) {
        return result.apply(null, null);
    }
    boolean[] nulls = null;
    long[] longs = null;
    RecordCursor cursor = recordSet.cursor();
    for (int position = 0; position < positions; position++) {
        checkState(cursor.advanceNextPosition(), "cursor has less values than expected");
        if (cursor.isNull(columnIndex)) {
            if (nulls == null) {
                nulls = new boolean[positions];
            }
            nulls[position] = true;
        }
        else {
            if (longs == null) {
                longs = new long[positions];
            }
            longs[position] = cursor.getLong(columnIndex);
        }
    }
    checkState(!cursor.advanceNextPosition(), "cursor has more values than expected");
    return result.apply(nulls, longs);
}
 

@Test
public void testCursor()
{
    ListBasedRecordSet recordSet = new ListBasedRecordSet(
            ImmutableList.of(
                    Arrays.asList("1", null, "3"),
                    Arrays.asList("ab", "c", null)),
            ImmutableList.of(BIGINT, VARCHAR));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, VARCHAR));
    RecordCursor cursor = recordSet.cursor();
    assertTrue(cursor.advanceNextPosition());
    assertEquals(cursor.getType(0), BIGINT);
    assertEquals(cursor.getType(1), VARCHAR);
    assertThrows(IndexOutOfBoundsException.class, () -> cursor.getLong(2));
    assertEquals(cursor.getLong(0), 1L);
    assertEquals(cursor.getSlice(1), Slices.utf8Slice("ab"));
    assertTrue(cursor.advanceNextPosition());
    assertTrue(cursor.isNull(0));
    assertEquals(cursor.getSlice(1), Slices.utf8Slice("c"));
    assertTrue(cursor.advanceNextPosition());
    assertEquals(cursor.getLong(0), 3L);
    assertTrue(cursor.isNull(1));
    assertFalse(cursor.advanceNextPosition());
    assertThrows(IndexOutOfBoundsException.class, () -> cursor.getLong(0));
}
 

private RecordCursor getCursor(JdbcTableHandle jdbcTableHandle, List<JdbcColumnHandle> columns, TupleDomain<ColumnHandle> domain)
{
    jdbcTableHandle = new JdbcTableHandle(
            jdbcTableHandle.getSchemaTableName(),
            jdbcTableHandle.getCatalogName(),
            jdbcTableHandle.getSchemaName(),
            jdbcTableHandle.getTableName(),
            domain,
            Optional.empty(),
            OptionalLong.empty(),
            Optional.empty());

    ConnectorSplitSource splits = jdbcClient.getSplits(SESSION, jdbcTableHandle);
    JdbcSplit split = (JdbcSplit) getOnlyElement(getFutureValue(splits.getNextBatch(NOT_PARTITIONED, 1000)).getSplits());

    ConnectorTransactionHandle transaction = new JdbcTransactionHandle();
    JdbcRecordSetProvider recordSetProvider = new JdbcRecordSetProvider(jdbcClient);
    RecordSet recordSet = recordSetProvider.getRecordSet(transaction, SESSION, split, jdbcTableHandle, columns);

    return recordSet.cursor();
}
 

private static RecordSet toRecordSet(ConnectorTransactionHandle sourceTransaction, SystemTable table, ConnectorSession session, TupleDomain<Integer> constraint)
{
    return new RecordSet()
    {
        private final List<Type> types = table.getTableMetadata().getColumns().stream()
                .map(ColumnMetadata::getType)
                .collect(toImmutableList());

        @Override
        public List<Type> getColumnTypes()
        {
            return types;
        }

        @Override
        public RecordCursor cursor()
        {
            return table.cursor(sourceTransaction, session, constraint);
        }
    };
}
 

public static PrestoThriftBlock fromIntBasedColumn(RecordSet recordSet, int columnIndex, int positions, BiFunction<boolean[], int[], PrestoThriftBlock> result)
{
    if (positions == 0) {
        return result.apply(null, null);
    }
    boolean[] nulls = null;
    int[] ints = null;
    RecordCursor cursor = recordSet.cursor();
    for (int position = 0; position < positions; position++) {
        checkState(cursor.advanceNextPosition(), "cursor has less values than expected");
        if (cursor.isNull(columnIndex)) {
            if (nulls == null) {
                nulls = new boolean[positions];
            }
            nulls[position] = true;
        }
        else {
            if (ints == null) {
                ints = new int[positions];
            }
            ints[position] = (int) cursor.getLong(columnIndex);
        }
    }
    checkState(!cursor.advanceNextPosition(), "cursor has more values than expected");
    return result.apply(nulls, ints);
}
 

@Override
public RecordCursor cursor()
{
    KuduScanner scanner = clientSession.createScanner(kuduSplit);
    Schema projectedSchema = scanner.getProjectionSchema();
    ImmutableMap.Builder<Integer, Integer> builder = ImmutableMap.builder();
    for (int i = 0; i < columns.size(); i++) {
        KuduColumnHandle handle = (KuduColumnHandle) columns.get(i);
        if (handle.isVirtualRowId()) {
            builder.put(i, ROW_ID_POSITION);
        }
        else {
            builder.put(i, projectedSchema.getColumnIndex(handle.getName()));
        }
    }

    return new KuduRecordCursor(scanner, getTable(), getColumnTypes(), builder.build());
}
 

public HiveReaderProjectionsAdaptingRecordCursor(RecordCursor delegate, ReaderProjectionsAdapter projectionsAdapter)
{
    this.delegate = requireNonNull(delegate, "delegate is null");
    requireNonNull(projectionsAdapter, "projectionsAdapter is null");

    this.channelMappings = new ChannelMapping[projectionsAdapter.getOutputToInputMapping().size()];
    projectionsAdapter.getOutputToInputMapping().toArray(channelMappings);

    this.outputTypes = new Type[projectionsAdapter.getOutputTypes().size()];
    projectionsAdapter.getOutputTypes().toArray(outputTypes);

    this.inputTypes = new Type[projectionsAdapter.getInputTypes().size()];
    projectionsAdapter.getInputTypes().toArray(inputTypes);

    this.baseTypes = new Type[outputTypes.length];
    for (int i = 0; i < baseTypes.length; i++) {
        Type type = inputTypes[channelMappings[i].getInputChannelIndex()];
        List<Integer> dereferences = channelMappings[i].getDereferenceSequence();
        for (int j = 0; j < dereferences.size(); j++) {
            type = type.getTypeParameters().get(dereferences.get(j));
        }
        baseTypes[i] = type;
    }
}
 

private Comparable<?> getPrestoValue(RecordCursor recordCursor, List<Column> columns, int columnId)
{
    if (recordCursor.isNull(columnId)) {
        return null;
    }

    Column column = columns.get(columnId);
    ColumnType.Base baseType = column.getType().getBase();
    switch (baseType) {
        case IDENTIFIER:
        case INTEGER:
        case DATE:
        case TIME:
        case DECIMAL:
            return recordCursor.getLong(columnId);
        case VARCHAR:
        case CHAR:
            return recordCursor.getSlice(columnId).toStringAscii();
    }
    throw new UnsupportedOperationException(format("Unsupported TPCDS base type [%s]", baseType));
}
 

private List<Long> readTimeStampsFrom(RecordSet recordSet)
{
    ImmutableList.Builder<Long> result = ImmutableList.builder();
    try (RecordCursor cursor = recordSet.cursor()) {
        while (cursor.advanceNextPosition()) {
            for (int i = 0; i < recordSet.getColumnTypes().size(); i++) {
                cursor.isNull(i);
            }
            if (cursor.isNull(0)) {
                return result.build();
            }
            assertTrue(recordSet.getColumnTypes().get(0) instanceof TimestampType);
            result.add(cursor.getLong(0));
        }
    }
    return result.build();
}
 

private static Object extractObject(RecordCursor cursor, int field, Type type)
{
    if (cursor.isNull(field)) {
        return null;
    }

    Class<?> javaType = type.getJavaType();
    if (javaType == boolean.class) {
        return cursor.getBoolean(field);
    }
    if (javaType == long.class) {
        return cursor.getLong(field);
    }
    if (javaType == double.class) {
        return cursor.getDouble(field);
    }
    if (javaType == Slice.class) {
        return cursor.getSlice(field).toStringUtf8();
    }
    throw new AssertionError("Unsupported type: " + type);
}
 

@Test
public void testCursorMixedOrder()
{
    RecordSet recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(
            new ExampleColumnHandle("value", BIGINT, 1),
            new ExampleColumnHandle("value", BIGINT, 1),
            new ExampleColumnHandle("text", createUnboundedVarcharType(), 0)));
    RecordCursor cursor = recordSet.cursor();

    Map<String, Long> data = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        assertEquals(cursor.getLong(0), cursor.getLong(1));
        data.put(cursor.getSlice(2).toStringUtf8(), cursor.getLong(0));
    }
    assertEquals(data, ImmutableMap.<String, Long>builder()
            .put("ten", 10L)
            .put("eleven", 11L)
            .put("twelve", 12L)
            .build());
}
 

@Test
public void testCursorSimple()
{
    RecordSet recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(
            new ExampleColumnHandle("text", createUnboundedVarcharType(), 0),
            new ExampleColumnHandle("value", BIGINT, 1)));
    RecordCursor cursor = recordSet.cursor();

    assertEquals(cursor.getType(0), createUnboundedVarcharType());
    assertEquals(cursor.getType(1), BIGINT);

    Map<String, Long> data = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        data.put(cursor.getSlice(0).toStringUtf8(), cursor.getLong(1));
        assertFalse(cursor.isNull(0));
        assertFalse(cursor.isNull(1));
    }
    assertEquals(data, ImmutableMap.<String, Long>builder()
            .put("ten", 10L)
            .put("eleven", 11L)
            .put("twelve", 12L)
            .build());
}
 

@Test
public void testGetRecordSet()
{
    ConnectorTableHandle tableHandle = new ExampleTableHandle("schema", "table");
    ExampleRecordSetProvider recordSetProvider = new ExampleRecordSetProvider();
    RecordSet recordSet = recordSetProvider.getRecordSet(ExampleTransactionHandle.INSTANCE, SESSION, new ExampleSplit(dataUri), tableHandle, ImmutableList.of(
            new ExampleColumnHandle("text", createUnboundedVarcharType(), 0),
            new ExampleColumnHandle("value", BIGINT, 1)));
    assertNotNull(recordSet, "recordSet is null");

    RecordCursor cursor = recordSet.cursor();
    assertNotNull(cursor, "cursor is null");

    Map<String, Long> data = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        data.put(cursor.getSlice(0).toStringUtf8(), cursor.getLong(1));
    }
    assertEquals(data, ImmutableMap.<String, Long>builder()
            .put("ten", 10L)
            .put("eleven", 11L)
            .put("twelve", 12L)
            .build());
}
 

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint)
{
    Builder table = InMemoryRecordSet.builder(transactionsTable);
    for (TransactionInfo info : transactionManager.getAllTransactionInfos()) {
        table.addRow(
                info.getTransactionId().toString(),
                info.getIsolationLevel().toString(),
                info.isReadOnly(),
                info.isAutoCommitContext(),
                info.getCreateTime().getMillis(),
                (long) info.getIdleTime().getValue(TimeUnit.SECONDS),
                info.getWrittenConnectorId().map(CatalogName::getCatalogName).orElse(null),
                createStringsBlock(info.getCatalogNames()));
    }
    return table.build().cursor();
}
 

private static IndexedTable indexTable(RecordSet recordSet, final List<String> outputColumns, List<String> keyColumns)
{
    List<Integer> keyPositions = keyColumns.stream()
            .map(columnName -> {
                int position = outputColumns.indexOf(columnName);
                checkState(position != -1);
                return position;
            })
            .collect(toImmutableList());

    ImmutableListMultimap.Builder<MaterializedTuple, MaterializedTuple> indexedValuesBuilder = ImmutableListMultimap.builder();

    List<Type> outputTypes = recordSet.getColumnTypes();
    List<Type> keyTypes = extractPositionValues(outputTypes, keyPositions);

    RecordCursor cursor = recordSet.cursor();
    while (cursor.advanceNextPosition()) {
        List<Object> values = extractValues(cursor, outputTypes);
        List<Object> keyValues = extractPositionValues(values, keyPositions);

        indexedValuesBuilder.put(new MaterializedTuple(keyValues), new MaterializedTuple(values));
    }

    return new IndexedTable(keyColumns, keyTypes, outputColumns, outputTypes, indexedValuesBuilder.build());
}
 

private static boolean fieldSetsEqual(RecordCursor cursor, List<Set<Field>> fieldSets)
{
    for (Set<Field> fieldSet : fieldSets) {
        if (!fieldsEquals(cursor, fieldSet)) {
            return false;
        }
    }
    return true;
}
 

private void generateProjectMethod(
        ClassDefinition classDefinition,
        CallSiteBinder callSiteBinder,
        CachedInstanceBinder cachedInstanceBinder,
        Map<LambdaDefinitionExpression, CompiledLambda> compiledLambdaMap,
        String methodName,
        RowExpression projection)
{
    Parameter session = arg("session", ConnectorSession.class);
    Parameter cursor = arg("cursor", RecordCursor.class);
    Parameter output = arg("output", BlockBuilder.class);
    MethodDefinition method = classDefinition.declareMethod(a(PUBLIC), methodName, type(void.class), session, cursor, output);

    method.comment("Projection: %s", projection.toString());

    Scope scope = method.getScope();
    Variable wasNullVariable = scope.declareVariable(type(boolean.class), "wasNull");

    RowExpressionCompiler compiler = new RowExpressionCompiler(
            callSiteBinder,
            cachedInstanceBinder,
            fieldReferenceCompiler(cursor),
            metadata,
            compiledLambdaMap);

    method.getBody()
            .comment("boolean wasNull = false;")
            .putVariable(wasNullVariable, false)
            .getVariable(output)
            .comment("evaluate projection: " + projection.toString())
            .append(compiler.compile(projection, scope))
            .append(generateWrite(callSiteBinder, scope, wasNullVariable, projection.getType()))
            .ret();
}
 

private ConcatRecordCursor(Iterator<RecordCursor> iterator, List<Type> types)
{
    // NOTE: this cursor implementation relies on the fact that the
    // cursor creation in the Iterable is lazy so DO NOT materialize this into
    // an ImmutableList
    this.iterator = requireNonNull(iterator, "iterator is null");
    this.types = ImmutableList.copyOf(requireNonNull(types, "types is null"));
}
 

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession connectorSession, TupleDomain<Integer> constraint)
{
    Session session = ((FullConnectorSession) connectorSession).getSession();
    Optional<String> catalogFilter = tryGetSingleVarcharValue(constraint, 0);
    Optional<String> schemaFilter = tryGetSingleVarcharValue(constraint, 1);
    Optional<String> tableFilter = tryGetSingleVarcharValue(constraint, 2);
    Optional<String> typeFilter = tryGetSingleVarcharValue(constraint, 3);

    boolean includeTables = emptyOrEquals(typeFilter, "TABLE");
    boolean includeViews = emptyOrEquals(typeFilter, "VIEW");
    Builder table = InMemoryRecordSet.builder(METADATA);

    if (!includeTables && !includeViews) {
        return table.build().cursor();
    }

    for (String catalog : listCatalogs(session, metadata, accessControl, catalogFilter).keySet()) {
        QualifiedTablePrefix prefix = tablePrefix(catalog, schemaFilter, tableFilter);

        Set<SchemaTableName> views = listViews(session, metadata, accessControl, prefix);
        for (SchemaTableName name : listTables(session, metadata, accessControl, prefix)) {
            boolean isView = views.contains(name);
            if ((includeTables && !isView) || (includeViews && isView)) {
                table.addRow(tableRow(catalog, name, isView ? "VIEW" : "TABLE"));
            }
        }
    }
    return table.build().cursor();
}
 

private static boolean fieldEquals(RecordCursor cursor, Field field1, Field field2)
{
    checkArgument(cursor.getType(field1.getField()).equals(cursor.getType(field2.getField())), "Should only be comparing fields of the same type");

    if (cursor.isNull(field1.getField()) || cursor.isNull(field2.getField())) {
        return false;
    }

    Class<?> javaType = cursor.getType(field1.getField()).getJavaType();
    try {
        if (javaType == long.class) {
            return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invokeExact(cursor.getLong(field1.getField()), cursor.getLong(field2.getField())));
        }
        if (javaType == double.class) {
            return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invokeExact(cursor.getDouble(field1.getField()), cursor.getDouble(field2.getField())));
        }
        if (javaType == boolean.class) {
            return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invokeExact(cursor.getBoolean(field1.getField()), cursor.getBoolean(field2.getField())));
        }
        if (javaType == Slice.class) {
            return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invokeExact(cursor.getSlice(field1.getField()), cursor.getSlice(field2.getField())));
        }
        return TRUE.equals((Boolean) field1.getEqualsMethodHandle().invoke(cursor.getObject(field1.getField()), cursor.getObject(field2.getField())));
    }
    catch (Throwable t) {
        throwIfUnchecked(t);
        throw new RuntimeException(t);
    }
}
 

private static boolean fieldsEquals(RecordCursor cursor, Set<Field> fields)
{
    if (fields.size() < 2) {
        return true; // Nothing to compare
    }
    Iterator<Field> fieldIterator = fields.iterator();
    Field firstField = fieldIterator.next();
    while (fieldIterator.hasNext()) {
        if (!fieldEquals(cursor, firstField, fieldIterator.next())) {
            return false;
        }
    }
    return true;
}
 

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint)
{
    return InMemoryRecordSet.builder(METADATA)
            .addRow("TABLE")
            .addRow("VIEW")
            .build().cursor();
}
 

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession connectorSession, TupleDomain<Integer> constraint)
{
    Session session = ((FullConnectorSession) connectorSession).getSession();
    Builder table = InMemoryRecordSet.builder(CATALOG_TABLE);
    for (Map.Entry<String, CatalogName> entry : listCatalogs(session, metadata, accessControl).entrySet()) {
        table.addRow(entry.getKey(), entry.getValue().toString());
    }
    return table.build().cursor();
}
 

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint)
{
    Builder table = InMemoryRecordSet.builder(NODES_TABLE);
    AllNodes allNodes = nodeManager.getAllNodes();
    addRows(table, allNodes.getActiveNodes(), ACTIVE);
    addRows(table, allNodes.getInactiveNodes(), INACTIVE);
    addRows(table, allNodes.getShuttingDownNodes(), SHUTTING_DOWN);
    return table.build().cursor();
}
 

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint)
{
    checkState(dispatchManager.isPresent(), "Query system table can return results only on coordinator");

    List<BasicQueryInfo> queries = dispatchManager.get().getQueries();
    queries = filterQueries(((FullConnectorSession) session).getSession().getIdentity(), queries, accessControl);

    Builder table = InMemoryRecordSet.builder(QUERY_TABLE);
    for (BasicQueryInfo queryInfo : queries) {
        Optional<QueryInfo> fullQueryInfo = dispatchManager.get().getFullQueryInfo(queryInfo.getQueryId());
        if (fullQueryInfo.isEmpty()) {
            continue;
        }
        QueryStats queryStats = fullQueryInfo.get().getQueryStats();
        table.addRow(
                queryInfo.getQueryId().toString(),
                queryInfo.getState().toString(),
                queryInfo.getSession().getUser(),
                queryInfo.getSession().getSource().orElse(null),
                queryInfo.getQuery(),
                queryInfo.getResourceGroupId().map(QuerySystemTable::resourceGroupIdToBlock).orElse(null),

                toMillis(queryStats.getQueuedTime()),
                toMillis(queryStats.getAnalysisTime()),
                toMillis(queryStats.getPlanningTime()),

                toTimeStamp(queryStats.getCreateTime()),
                toTimeStamp(queryStats.getExecutionStartTime()),
                toTimeStamp(queryStats.getLastHeartbeat()),
                toTimeStamp(queryStats.getEndTime()),

                Optional.ofNullable(queryInfo.getErrorType()).map(Enum::name).orElse(null),
                Optional.ofNullable(queryInfo.getErrorCode()).map(ErrorCode::getName).orElse(null));
    }
    return table.build().cursor();
}
 

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession connectorSession, TupleDomain<Integer> constraint)
{
    Session session = ((FullConnectorSession) connectorSession).getSession();
    Optional<String> catalogFilter = tryGetSingleVarcharValue(constraint, 1);

    Builder table = InMemoryRecordSet.builder(METADATA);
    for (String catalog : listCatalogs(session, metadata, accessControl, catalogFilter).keySet()) {
        for (String schema : listSchemas(session, metadata, accessControl, catalog)) {
            table.addRow(schema, catalog);
        }
    }
    return table.build().cursor();
}
 

@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession connectorSession, TupleDomain<Integer> constraint)
{
    Session session = ((FullConnectorSession) connectorSession).getSession();
    Builder table = InMemoryRecordSet.builder(METADATA);
    for (String name : listCatalogs(session, metadata, accessControl).keySet()) {
        table.addRow(name);
    }
    return table.build().cursor();
}
 

private static List<Object> extractValues(RecordCursor cursor, List<Type> types)
{
    List<Object> list = new ArrayList<>(types.size());
    for (int i = 0; i < types.size(); i++) {
        list.add(extractObject(cursor, i, types.get(i)));
    }
    return list;
}
 

private static void assertRecordCursor(List<Type> types, Iterator<?>[] valuesByField, RecordCursor cursor)
{
    while (cursor.advanceNextPosition()) {
        for (int field = 0; field < types.size(); field++) {
            assertTrue(valuesByField[field].hasNext());
            Object expected = valuesByField[field].next();
            Object actual = getActualCursorValue(cursor, types.get(field), field);
            assertEquals(actual, expected);
        }
    }
}