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

Example 11 with RecordCursor

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

the class TableStatisticsRecorder method recordStatistics.

public TableStatisticsData recordStatistics(Table table, double scaleFactor) {
    Session session = Session.getDefaultSession().withScale(scaleFactor).withParallelism(1).withNoSexism(false);
    List<Column> columns = ImmutableList.copyOf(table.getColumns());
    RecordCursor recordCursor = new TpcdsRecordSet(Results.constructResults(table, session), columns).cursor();
    List<ColumnStatisticsRecorder> statisticsRecorders = createStatisticsRecorders(columns);
    long rowCount = 0;
    while (recordCursor.advanceNextPosition()) {
        rowCount++;
        for (int columnId = 0; columnId < columns.size(); columnId++) {
            Comparable<?> value = getTrinoValue(recordCursor, columns, columnId);
            statisticsRecorders.get(columnId).record(value);
        }
    }
    Map<String, ColumnStatisticsData> columnSampleStatistics = IntStream.range(0, columns.size()).boxed().collect(toImmutableMap(i -> columns.get(i).getName(), i -> statisticsRecorders.get(i).getRecording()));
    return new TableStatisticsData(rowCount, columnSampleStatistics);
}

Example 12 with RecordCursor

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

the class H2QueryRunner method insertRows.

private static void insertRows(ConnectorTableMetadata tableMetadata, Handle handle, RecordSet data) {
    List<ColumnMetadata> columns = tableMetadata.getColumns().stream().filter(columnMetadata -> !columnMetadata.isHidden()).collect(toImmutableList());
    String vars = Joiner.on(',').join(nCopies(columns.size(), "?"));
    String sql = format("INSERT INTO %s VALUES (%s)", tableMetadata.getTable().getTableName(), vars);
    RecordCursor cursor = data.cursor();
    while (true) {
        // insert 1000 rows at a time
        PreparedBatch batch = handle.prepareBatch(sql);
        for (int row = 0; row < 1000; row++) {
            if (!cursor.advanceNextPosition()) {
                if (batch.size() > 0) {
                    batch.execute();
                }
                return;
            }
            for (int column = 0; column < columns.size(); column++) {
                Type type = columns.get(column).getType();
                if (BOOLEAN.equals(type)) {
                    batch.bind(column, cursor.getBoolean(column));
                } else if (BIGINT.equals(type)) {
                    batch.bind(column, cursor.getLong(column));
                } else if (INTEGER.equals(type)) {
                    batch.bind(column, toIntExact(cursor.getLong(column)));
                } else if (DOUBLE.equals(type)) {
                    batch.bind(column, cursor.getDouble(column));
                } else if (type instanceof VarcharType) {
                    batch.bind(column, cursor.getSlice(column).toStringUtf8());
                } else if (DATE.equals(type)) {
                    long millisUtc = TimeUnit.DAYS.toMillis(cursor.getLong(column));
                    // H2 expects dates in to be millis at midnight in the JVM timezone
                    long localMillis = DateTimeZone.UTC.getMillisKeepLocal(DateTimeZone.getDefault(), millisUtc);
                    batch.bind(column, new Date(localMillis));
                } else {
                    throw new IllegalArgumentException("Unsupported type " + type);
                }
            }
            batch.add();
        }
        batch.execute();
    }
}

Example 13 with RecordCursor

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

the class TestListBasedRecordSet method testCursor.

@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);
    assertThatThrownBy(() -> cursor.getLong(2)).isInstanceOf(IndexOutOfBoundsException.class).hasMessage("Index 2 out of bounds for length 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());
    assertThatThrownBy(() -> cursor.getLong(0)).isInstanceOf(IndexOutOfBoundsException.class).hasMessage("Index 3 out of bounds for length 3");
}

Example 14 with RecordCursor

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

the class TestListBasedRecordSet method testEmptyCursor.

@Test
public void testEmptyCursor() {
    ListBasedRecordSet recordSet = new ListBasedRecordSet(ImmutableList.of(), ImmutableList.of(BIGINT, INTEGER));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, INTEGER));
    RecordCursor cursor = recordSet.cursor();
    assertFalse(cursor.advanceNextPosition());
}

Example 15 with RecordCursor

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

the class TrinoThriftTypeUtils method fromIntBasedColumn.

public static TrinoThriftBlock fromIntBasedColumn(RecordSet recordSet, int columnIndex, int positions, BiFunction<boolean[], int[], TrinoThriftBlock> 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);
}