use of io.prestosql.spi.connector.RecordCursor in project hetu-core by openlookeng.
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, (int) 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();
}
}
use of io.prestosql.spi.connector.RecordCursor in project hetu-core by openlookeng.
the class TableCommentSystemTable method cursor.
@Override
public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession connectorSession, TupleDomain<Integer> constraint) {
Optional<String> catalogFilter = stringFilter(constraint, 0);
Optional<String> schemaFilter = stringFilter(constraint, 1);
Optional<String> tableFilter = stringFilter(constraint, 2);
Session session = toSession(transactionHandle, connectorSession);
Builder table = InMemoryRecordSet.builder(COMMENT_TABLE);
for (String catalog : filter(listCatalogs(session, metadata, accessControl).keySet(), catalogFilter)) {
QualifiedTablePrefix prefix = tablePrefix(catalog, schemaFilter, tableFilter);
Set<SchemaTableName> names = ImmutableSet.of();
try {
names = listTables(session, metadata, accessControl, prefix);
} catch (PrestoException e) {
// listTables throws an exception if cannot connect the database
LOG.debug(e, "Failed to get tables for catalog: %s", catalog);
}
for (SchemaTableName name : names) {
QualifiedObjectName tableName = new QualifiedObjectName(prefix.getCatalogName(), name.getSchemaName(), name.getTableName());
Optional<String> comment = Optional.empty();
try {
comment = metadata.getTableHandle(session, tableName).map(handle -> metadata.getTableMetadata(session, handle)).map(metadata -> metadata.getMetadata().getComment()).get();
} catch (PrestoException e) {
// getTableHandle may throw an exception (e.g. Cassandra connector doesn't allow case insensitive column names)
LOG.debug(e, "Failed to get metadata for table: %s", name);
}
table.addRow(prefix.getCatalogName(), name.getSchemaName(), name.getTableName(), comment.orElse(null));
}
}
return table.build().cursor();
}
use of io.prestosql.spi.connector.RecordCursor in project hetu-core by openlookeng.
the class TestJdbcRecordSet method testCursorSimple.
@Test
public void testCursorSimple() {
RecordSet recordSet = new JdbcRecordSet(jdbcClient, SESSION, split, table, ImmutableList.of(columnHandles.get("text"), columnHandles.get("text_short"), columnHandles.get("value")));
try (RecordCursor cursor = recordSet.cursor()) {
assertEquals(cursor.getType(0), VARCHAR);
assertEquals(cursor.getType(1), createVarcharType(32));
assertEquals(cursor.getType(2), BIGINT);
Map<String, Long> data = new LinkedHashMap<>();
while (cursor.advanceNextPosition()) {
data.put(cursor.getSlice(0).toStringUtf8(), cursor.getLong(2));
assertEquals(cursor.getSlice(0), cursor.getSlice(1));
assertFalse(cursor.isNull(0));
assertFalse(cursor.isNull(1));
assertFalse(cursor.isNull(2));
}
assertEquals(data, ImmutableMap.<String, Long>builder().put("one", 1L).put("two", 2L).put("three", 3L).put("ten", 10L).put("eleven", 11L).put("twelve", 12L).build());
}
}
use of io.prestosql.spi.connector.RecordCursor in project hetu-core by openlookeng.
the class TestJdbcRecordSet method testIdempotentClose.
@Test
public void testIdempotentClose() {
RecordSet recordSet = new JdbcRecordSet(jdbcClient, SESSION, split, table, ImmutableList.of(columnHandles.get("value"), columnHandles.get("value"), columnHandles.get("text")));
RecordCursor cursor = recordSet.cursor();
cursor.close();
cursor.close();
}
use of io.prestosql.spi.connector.RecordCursor in project hetu-core by openlookeng.
the class AbstractTestHiveFileFormats method checkCursor.
protected void checkCursor(RecordCursor cursor, List<TestColumn> testColumns, int rowCount) {
List<Type> types = testColumns.stream().map(column -> column.getObjectInspector().getTypeName()).map(type -> HiveType.valueOf(type).getType(TYPE_MANAGER)).collect(toImmutableList());
Map<Type, MethodHandle> distinctFromOperators = types.stream().distinct().collect(toImmutableMap(identity(), HiveTestUtils::distinctFromOperator));
for (int row = 0; row < rowCount; row++) {
assertTrue(cursor.advanceNextPosition());
for (int i = 0, testColumnsSize = testColumns.size(); i < testColumnsSize; i++) {
TestColumn testColumn = testColumns.get(i);
Type type = types.get(i);
Object fieldFromCursor = getFieldFromCursor(cursor, type, i);
if (fieldFromCursor == null) {
assertEquals(null, testColumn.getExpectedValue(), "Expected null for column " + testColumn.getName());
} else if (type instanceof DecimalType) {
DecimalType decimalType = (DecimalType) type;
fieldFromCursor = new BigDecimal((BigInteger) fieldFromCursor, decimalType.getScale());
assertEquals(fieldFromCursor, testColumn.getExpectedValue(), "Wrong value for column " + testColumn.getName());
} else if (testColumn.getObjectInspector().getTypeName().equals("float")) {
assertEquals((float) fieldFromCursor, (float) testColumn.getExpectedValue(), (float) EPSILON);
} else if (testColumn.getObjectInspector().getTypeName().equals("double")) {
assertEquals((double) fieldFromCursor, (double) testColumn.getExpectedValue(), EPSILON);
} else if (testColumn.getObjectInspector().getTypeName().equals("tinyint")) {
assertEquals(((Number) fieldFromCursor).byteValue(), testColumn.getExpectedValue());
} else if (testColumn.getObjectInspector().getTypeName().equals("smallint")) {
assertEquals(((Number) fieldFromCursor).shortValue(), testColumn.getExpectedValue());
} else if (testColumn.getObjectInspector().getTypeName().equals("int")) {
assertEquals(((Number) fieldFromCursor).intValue(), testColumn.getExpectedValue());
} else if (testColumn.getObjectInspector().getCategory() == Category.PRIMITIVE) {
assertEquals(fieldFromCursor, testColumn.getExpectedValue(), "Wrong value for column " + testColumn.getName());
} else {
Block expected = (Block) testColumn.getExpectedValue();
Block actual = (Block) fieldFromCursor;
boolean distinct = isDistinctFrom(distinctFromOperators.get(type), expected, actual);
assertFalse(distinct, "Wrong value for column: " + testColumn.getName());
}
}
}
assertFalse(cursor.advanceNextPosition());
}
Aggregations