Examples with RecordCursor io.trino.spi.connector.RecordCursor used on opensource projects

Example 16 with RecordCursor

use of io.trino.spi.connector.RecordCursor in project trino by trinodb.

the class TestCassandraConnector method testGetTupleType.

@Test
public void testGetTupleType() {
    // TODO add test with nested tuple types
    ConnectorTableHandle tableHandle = getTableHandle(tableTuple);
    ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(SESSION, tableHandle);
    List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(SESSION, tableHandle).values());
    Map<String, Integer> columnIndex = indexColumns(columnHandles);
    ConnectorTransactionHandle transaction = CassandraTransactionHandle.INSTANCE;
    List<ConnectorSplit> splits = getAllSplits(splitManager.getSplits(transaction, SESSION, tableHandle, UNGROUPED_SCHEDULING, DynamicFilter.EMPTY));
    long rowNumber = 0;
    for (ConnectorSplit split : splits) {
        CassandraSplit cassandraSplit = (CassandraSplit) split;
        long completedBytes = 0;
        try (RecordCursor cursor = recordSetProvider.getRecordSet(transaction, SESSION, cassandraSplit, tableHandle, columnHandles).cursor()) {
            while (cursor.advanceNextPosition()) {
                try {
                    assertReadFields(cursor, tableMetadata.getColumns());
                } catch (RuntimeException e) {
                    throw new RuntimeException("row " + rowNumber, e);
                }
                rowNumber++;
                String keyValue = cursor.getSlice(columnIndex.get("key")).toStringUtf8();
                assertEquals(keyValue, Long.toString(rowNumber));
                SingleRowBlock tupleValueBlock = (SingleRowBlock) cursor.getObject(columnIndex.get("typetuple"));
                assertThat(tupleValueBlock.getPositionCount()).isEqualTo(3);
                CassandraColumnHandle tupleColumnHandle = (CassandraColumnHandle) columnHandles.get(columnIndex.get("typetuple"));
                List<CassandraType> tupleArgumentTypes = tupleColumnHandle.getCassandraType().getArgumentTypes();
                assertThat(tupleArgumentTypes.get(0).getTrinoType().getLong(tupleValueBlock, 0)).isEqualTo(rowNumber);
                assertThat(tupleArgumentTypes.get(1).getTrinoType().getSlice(tupleValueBlock, 1).toStringUtf8()).isEqualTo("text-" + rowNumber);
                assertThat(tupleArgumentTypes.get(2).getTrinoType().getLong(tupleValueBlock, 2)).isEqualTo(Float.floatToRawIntBits(1.11f * rowNumber));
                long newCompletedBytes = cursor.getCompletedBytes();
                assertTrue(newCompletedBytes >= completedBytes);
                completedBytes = newCompletedBytes;
            }
        }
    }
    assertEquals(rowNumber, 2);
}

Example 17 with RecordCursor

use of io.trino.spi.connector.RecordCursor in project trino by trinodb.

the class TestCassandraConnector method testGetUserDefinedType.

@Test
public void testGetUserDefinedType() {
    ConnectorTableHandle tableHandle = getTableHandle(tableUdt);
    ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(SESSION, tableHandle);
    List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(SESSION, tableHandle).values());
    Map<String, Integer> columnIndex = indexColumns(columnHandles);
    ConnectorTransactionHandle transaction = CassandraTransactionHandle.INSTANCE;
    tableHandle = metadata.applyFilter(SESSION, tableHandle, Constraint.alwaysTrue()).get().getHandle();
    List<ConnectorSplit> splits = getAllSplits(splitManager.getSplits(transaction, SESSION, tableHandle, UNGROUPED_SCHEDULING, DynamicFilter.EMPTY));
    long rowNumber = 0;
    for (ConnectorSplit split : splits) {
        CassandraSplit cassandraSplit = (CassandraSplit) split;
        long completedBytes = 0;
        try (RecordCursor cursor = recordSetProvider.getRecordSet(transaction, SESSION, cassandraSplit, tableHandle, columnHandles).cursor()) {
            while (cursor.advanceNextPosition()) {
                try {
                    assertReadFields(cursor, tableMetadata.getColumns());
                } catch (RuntimeException e) {
                    throw new RuntimeException("row " + rowNumber, e);
                }
                rowNumber++;
                String keyValue = cursor.getSlice(columnIndex.get("key")).toStringUtf8();
                SingleRowBlock udtValue = (SingleRowBlock) cursor.getObject(columnIndex.get("typeudt"));
                assertEquals(keyValue, "key");
                assertEquals(VARCHAR.getSlice(udtValue, 0).toStringUtf8(), "text");
                assertEquals(trinoUuidToJavaUuid(UUID.getSlice(udtValue, 1)).toString(), "01234567-0123-0123-0123-0123456789ab");
                assertEquals(INTEGER.getLong(udtValue, 2), -2147483648);
                assertEquals(BIGINT.getLong(udtValue, 3), -9223372036854775808L);
                assertEquals(VARBINARY.getSlice(udtValue, 4).toStringUtf8(), "01234");
                assertEquals(TIMESTAMP.getLong(udtValue, 5), 117964800000L);
                assertEquals(VARCHAR.getSlice(udtValue, 6).toStringUtf8(), "ansi");
                assertTrue(BOOLEAN.getBoolean(udtValue, 7));
                assertEquals(DOUBLE.getDouble(udtValue, 8), 99999999999999997748809823456034029568D);
                assertEquals(DOUBLE.getDouble(udtValue, 9), 4.9407e-324);
                assertEquals(REAL.getObjectValue(SESSION, udtValue, 10), 1.4E-45f);
                assertEquals(VARCHAR.getSlice(udtValue, 11).toStringUtf8(), "0.0.0.0");
                assertEquals(VARCHAR.getSlice(udtValue, 12).toStringUtf8(), "varchar");
                assertEquals(VARCHAR.getSlice(udtValue, 13).toStringUtf8(), "-9223372036854775808");
                assertEquals(trinoUuidToJavaUuid(UUID.getSlice(udtValue, 14)).toString(), "d2177dd0-eaa2-11de-a572-001b779c76e3");
                assertEquals(VARCHAR.getSlice(udtValue, 15).toStringUtf8(), "[\"list\"]");
                assertEquals(VARCHAR.getSlice(udtValue, 16).toStringUtf8(), "{\"map\":1}");
                assertEquals(VARCHAR.getSlice(udtValue, 17).toStringUtf8(), "[true]");
                SingleRowBlock tupleValueBlock = (SingleRowBlock) udtValue.getObject(18, Block.class);
                assertThat(tupleValueBlock.getPositionCount()).isEqualTo(1);
                assertThat(INTEGER.getLong(tupleValueBlock, 0)).isEqualTo(123);
                SingleRowBlock udtValueBlock = (SingleRowBlock) udtValue.getObject(19, Block.class);
                assertThat(udtValueBlock.getPositionCount()).isEqualTo(1);
                assertThat(INTEGER.getLong(udtValueBlock, 0)).isEqualTo(999);
                long newCompletedBytes = cursor.getCompletedBytes();
                assertTrue(newCompletedBytes >= completedBytes);
                completedBytes = newCompletedBytes;
            }
        }
    }
    assertEquals(rowNumber, 1);
}

Example 18 with RecordCursor

use of io.trino.spi.connector.RecordCursor in project trino by trinodb.

the class TestCassandraConnector method testGetRecords.

@Test
public void testGetRecords() {
    ConnectorTableHandle tableHandle = getTableHandle(table);
    ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(SESSION, tableHandle);
    List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(SESSION, tableHandle).values());
    Map<String, Integer> columnIndex = indexColumns(columnHandles);
    ConnectorTransactionHandle transaction = CassandraTransactionHandle.INSTANCE;
    tableHandle = metadata.applyFilter(SESSION, tableHandle, Constraint.alwaysTrue()).get().getHandle();
    List<ConnectorSplit> splits = getAllSplits(splitManager.getSplits(transaction, SESSION, tableHandle, UNGROUPED_SCHEDULING, DynamicFilter.EMPTY));
    long rowNumber = 0;
    for (ConnectorSplit split : splits) {
        CassandraSplit cassandraSplit = (CassandraSplit) split;
        long completedBytes = 0;
        try (RecordCursor cursor = recordSetProvider.getRecordSet(transaction, SESSION, cassandraSplit, tableHandle, columnHandles).cursor()) {
            while (cursor.advanceNextPosition()) {
                try {
                    assertReadFields(cursor, tableMetadata.getColumns());
                } catch (RuntimeException e) {
                    throw new RuntimeException("row " + rowNumber, e);
                }
                rowNumber++;
                String keyValue = cursor.getSlice(columnIndex.get("key")).toStringUtf8();
                assertTrue(keyValue.startsWith("key "));
                int rowId = Integer.parseInt(keyValue.substring(4));
                assertEquals(keyValue, "key " + rowId);
                assertEquals(Bytes.toHexString(cursor.getSlice(columnIndex.get("typebytes")).getBytes()), format("0x%08X", rowId));
                // VARINT is returned as a string
                assertEquals(cursor.getSlice(columnIndex.get("typeinteger")).toStringUtf8(), String.valueOf(rowId));
                assertEquals(cursor.getLong(columnIndex.get("typelong")), 1000 + rowId);
                assertEquals(trinoUuidToJavaUuid(cursor.getSlice(columnIndex.get("typeuuid"))).toString(), format("00000000-0000-0000-0000-%012d", rowId));
                assertEquals(cursor.getLong(columnIndex.get("typetimestamp")), packDateTimeWithZone(DATE.getTime(), UTC_KEY));
                long newCompletedBytes = cursor.getCompletedBytes();
                assertTrue(newCompletedBytes >= completedBytes);
                completedBytes = newCompletedBytes;
            }
        }
    }
    assertEquals(rowNumber, 9);
}

Example 19 with RecordCursor

use of io.trino.spi.connector.RecordCursor in project trino by trinodb.

the class S3SelectRecordCursorProvider method createRecordCursor.

@Override
public Optional<ReaderRecordCursorWithProjections> createRecordCursor(Configuration configuration, ConnectorSession session, Path path, long start, long length, long fileSize, Properties schema, List<HiveColumnHandle> columns, TupleDomain<HiveColumnHandle> effectivePredicate, TypeManager typeManager, boolean s3SelectPushdownEnabled) {
    if (!s3SelectPushdownEnabled) {
        return Optional.empty();
    }
    try {
        this.hdfsEnvironment.getFileSystem(session.getIdentity(), path, configuration);
    } catch (IOException e) {
        throw new TrinoException(HIVE_FILESYSTEM_ERROR, "Failed getting FileSystem: " + path, e);
    }
    Optional<ReaderColumns> projectedReaderColumns = projectBaseColumns(columns);
    // Ignore predicates on partial columns for now.
    effectivePredicate = effectivePredicate.filter((column, domain) -> column.isBaseColumn());
    String serdeName = getDeserializerClassName(schema);
    if (CSV_SERDES.contains(serdeName)) {
        List<HiveColumnHandle> readerColumns = projectedReaderColumns.map(ReaderColumns::get).map(readColumns -> readColumns.stream().map(HiveColumnHandle.class::cast).collect(toUnmodifiableList())).orElse(columns);
        IonSqlQueryBuilder queryBuilder = new IonSqlQueryBuilder(typeManager);
        String ionSqlQuery = queryBuilder.buildSql(readerColumns, effectivePredicate);
        S3SelectLineRecordReader recordReader = new S3SelectCsvRecordReader(configuration, path, start, length, schema, ionSqlQuery, s3ClientFactory);
        RecordCursor cursor = new S3SelectRecordCursor<>(configuration, path, recordReader, length, schema, readerColumns);
        return Optional.of(new ReaderRecordCursorWithProjections(cursor, projectedReaderColumns));
    }
    // unsupported serdes
    return Optional.empty();
}

Example 20 with RecordCursor

use of io.trino.spi.connector.RecordCursor in project trino by trinodb.

the class PartitionTable method buildRecordCursor.

private RecordCursor buildRecordCursor(Map<StructLikeWrapper, IcebergStatistics> partitionStatistics, List<PartitionField> partitionFields) {
    List<Type> partitionTypes = partitionTypes(partitionFields);
    List<? extends Class<?>> partitionColumnClass = partitionTypes.stream().map(type -> type.typeId().javaClass()).collect(toImmutableList());
    ImmutableList.Builder<List<Object>> records = ImmutableList.builder();
    for (Map.Entry<StructLikeWrapper, IcebergStatistics> partitionEntry : partitionStatistics.entrySet()) {
        StructLikeWrapper partitionStruct = partitionEntry.getKey();
        IcebergStatistics icebergStatistics = partitionEntry.getValue();
        List<Object> row = new ArrayList<>();
        // add data for partition columns
        partitionColumnType.ifPresent(partitionColumnType -> {
            BlockBuilder partitionRowBlockBuilder = partitionColumnType.createBlockBuilder(null, 1);
            BlockBuilder partitionBlockBuilder = partitionRowBlockBuilder.beginBlockEntry();
            for (int i = 0; i < partitionColumnTypes.size(); i++) {
                io.trino.spi.type.Type trinoType = partitionColumnType.getFields().get(i).getType();
                Object value = convertIcebergValueToTrino(partitionTypes.get(i), partitionStruct.get().get(i, partitionColumnClass.get(i)));
                writeNativeValue(trinoType, partitionBlockBuilder, value);
            }
            partitionRowBlockBuilder.closeEntry();
            row.add(partitionColumnType.getObject(partitionRowBlockBuilder, 0));
        });
        // add the top level metrics.
        row.add(icebergStatistics.getRecordCount());
        row.add(icebergStatistics.getFileCount());
        row.add(icebergStatistics.getSize());
        // add column level metrics
        dataColumnType.ifPresent(dataColumnType -> {
            BlockBuilder dataRowBlockBuilder = dataColumnType.createBlockBuilder(null, 1);
            BlockBuilder dataBlockBuilder = dataRowBlockBuilder.beginBlockEntry();
            for (int i = 0; i < columnMetricTypes.size(); i++) {
                Integer fieldId = nonPartitionPrimitiveColumns.get(i).fieldId();
                Object min = icebergStatistics.getMinValues().get(fieldId);
                Object max = icebergStatistics.getMaxValues().get(fieldId);
                Long nullCount = icebergStatistics.getNullCounts().get(fieldId);
                if (min == null && max == null && nullCount == null) {
                    row.add(null);
                    return;
                }
                RowType columnMetricType = columnMetricTypes.get(i);
                columnMetricType.writeObject(dataBlockBuilder, getColumnMetricBlock(columnMetricType, min, max, nullCount));
            }
            dataRowBlockBuilder.closeEntry();
            row.add(dataColumnType.getObject(dataRowBlockBuilder, 0));
        });
        records.add(row);
    }
    return new InMemoryRecordSet(resultTypes, records.build()).cursor();
}