@Test
public void getTableMetadata() {
// known table
ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(SESSION, tableHandle);
assertEquals(tableMetadata.getTable(), new SchemaTableName("example", "numbers"));
assertEquals(
tableMetadata.getColumns(),
ImmutableList.of(
new ColumnMetadata("text", VARCHAR, false),
new ColumnMetadata("value", BIGINT, false)));
// escaping name patterns
JdbcTableHandle specialTableHandle =
metadata.getTableHandle(SESSION, new SchemaTableName("exa_ple", "num_ers"));
ConnectorTableMetadata specialTableMetadata =
metadata.getTableMetadata(SESSION, specialTableHandle);
assertEquals(specialTableMetadata.getTable(), new SchemaTableName("exa_ple", "num_ers"));
assertEquals(
specialTableMetadata.getColumns(),
ImmutableList.of(
new ColumnMetadata("te_t", VARCHAR, false),
new ColumnMetadata("va%ue", BIGINT, false)));
// unknown tables should produce null
unknownTableMetadata(
new JdbcTableHandle(
CONNECTOR_ID, new SchemaTableName("u", "numbers"), null, "unknown", "unknown"));
unknownTableMetadata(
new JdbcTableHandle(
CONNECTOR_ID, new SchemaTableName("example", "numbers"), null, "example", "unknown"));
unknownTableMetadata(
new JdbcTableHandle(
CONNECTOR_ID, new SchemaTableName("example", "numbers"), null, "unknown", "numbers"));
}
@SuppressWarnings({"ValueOfIncrementOrDecrementUsed", "UnusedAssignment"})
@Test
public void testGetTableSchema() throws Exception {
ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(getTableHandle(table));
Map<String, ColumnMetadata> map = uniqueIndex(tableMetadata.getColumns(), columnNameGetter());
int i = 0;
assertPrimitiveField(map, i++, "t_string", VARCHAR, false);
assertPrimitiveField(map, i++, "t_tinyint", BIGINT, false);
assertPrimitiveField(map, i++, "t_smallint", BIGINT, false);
assertPrimitiveField(map, i++, "t_int", BIGINT, false);
assertPrimitiveField(map, i++, "t_bigint", BIGINT, false);
assertPrimitiveField(map, i++, "t_float", DOUBLE, false);
assertPrimitiveField(map, i++, "t_double", DOUBLE, false);
assertPrimitiveField(map, i++, "t_map", VARCHAR, false); // Currently mapped as a string
assertPrimitiveField(map, i++, "t_boolean", BOOLEAN, false);
assertPrimitiveField(map, i++, "t_timestamp", TIMESTAMP, false);
assertPrimitiveField(map, i++, "t_binary", VARBINARY, false);
assertPrimitiveField(
map, i++, "t_array_string", VARCHAR, false); // Currently mapped as a string
assertPrimitiveField(map, i++, "t_complex", VARCHAR, false); // Currently mapped as a string
assertPrimitiveField(map, i++, "ds", VARCHAR, true);
assertPrimitiveField(map, i++, "file_format", VARCHAR, true);
assertPrimitiveField(map, i++, "dummy", BIGINT, true);
}
@Override
public void createTable(ConnectorSession session, ConnectorTableMetadata tableMetadata) {
checkArgument(!isNullOrEmpty(tableMetadata.getOwner()), "Table owner is null or empty");
SchemaTableName schemaTableName = tableMetadata.getTable();
String schemaName = schemaTableName.getSchemaName();
String tableName = schemaTableName.getTableName();
ImmutableList.Builder<String> columnNames = ImmutableList.builder();
ImmutableList.Builder<Type> columnTypes = ImmutableList.builder();
buildColumnInfo(tableMetadata, columnNames, columnTypes);
ImmutableList.Builder<FieldSchema> partitionKeys = ImmutableList.builder();
ImmutableList.Builder<FieldSchema> columns = ImmutableList.builder();
List<String> names = columnNames.build();
List<String> typeNames =
columnTypes
.build()
.stream()
.map(HiveType::toHiveType)
.map(HiveType::getHiveTypeName)
.collect(toList());
for (int i = 0; i < names.size(); i++) {
if (tableMetadata.getColumns().get(i).isPartitionKey()) {
partitionKeys.add(new FieldSchema(names.get(i), typeNames.get(i), null));
} else {
columns.add(new FieldSchema(names.get(i), typeNames.get(i), null));
}
}
Path targetPath = getTargetPath(schemaName, tableName, schemaTableName);
HiveStorageFormat hiveStorageFormat = getHiveStorageFormat(session, this.hiveStorageFormat);
SerDeInfo serdeInfo = new SerDeInfo();
serdeInfo.setName(tableName);
serdeInfo.setSerializationLib(hiveStorageFormat.getSerDe());
StorageDescriptor sd = new StorageDescriptor();
sd.setLocation(targetPath.toString());
sd.setCols(columns.build());
sd.setSerdeInfo(serdeInfo);
sd.setInputFormat(hiveStorageFormat.getInputFormat());
sd.setOutputFormat(hiveStorageFormat.getOutputFormat());
Table table = new Table();
table.setDbName(schemaName);
table.setTableName(tableName);
table.setOwner(tableMetadata.getOwner());
table.setTableType(TableType.MANAGED_TABLE.toString());
String tableComment = "Created by Presto";
table.setParameters(ImmutableMap.of("comment", tableComment));
table.setPartitionKeys(partitionKeys.build());
table.setSd(sd);
metastore.createTable(table);
}
@Test
public void testGetTableSchemaOfflinePartition() throws Exception {
ConnectorTableHandle tableHandle = getTableHandle(tableOfflinePartition);
ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(tableHandle);
Map<String, ColumnMetadata> map = uniqueIndex(tableMetadata.getColumns(), columnNameGetter());
assertPrimitiveField(map, 0, "t_string", VARCHAR, false);
}
public static Map<String, ColumnHandle> toSystemColumnHandles(
ConnectorId connectorId, ConnectorTableMetadata tableMetadata) {
ImmutableMap.Builder<String, ColumnHandle> columnHandles = ImmutableMap.builder();
for (ColumnMetadata columnMetadata : tableMetadata.getColumns()) {
columnHandles.put(
columnMetadata.getName(), new SystemColumnHandle(connectorId, columnMetadata.getName()));
}
return columnHandles.build();
}
@Override
public Map<SchemaTableName, List<ColumnMetadata>> listTableColumns(SchemaTablePrefix prefix) {
checkNotNull(prefix, "prefix is null");
ImmutableMap.Builder<SchemaTableName, List<ColumnMetadata>> columns = ImmutableMap.builder();
for (SchemaTableName tableName : listTables(prefix)) {
ConnectorTableMetadata tableMetadata = getTableMetadata(tableName);
// table can disappear during listing operation
if (tableMetadata != null) {
columns.put(tableName, tableMetadata.getColumns());
}
}
return columns.build();
}
private static void insertRows(
ConnectorTableMetadata tableMetadata, Handle handle, RecordSet data) {
List<ColumnMetadata> columns =
ImmutableList.copyOf(
Iterables.filter(
tableMetadata.getColumns(),
new Predicate<ColumnMetadata>() {
@Override
public boolean apply(ColumnMetadata columnMetadata) {
return !columnMetadata.isHidden();
}
}));
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()) {
batch.execute();
return;
}
PreparedBatchPart part = batch.add();
for (int column = 0; column < columns.size(); column++) {
Type type = columns.get(column).getType();
if (BOOLEAN.equals(type)) {
part.bind(column, cursor.getBoolean(column));
} else if (BIGINT.equals(type)) {
part.bind(column, cursor.getLong(column));
} else if (DOUBLE.equals(type)) {
part.bind(column, cursor.getDouble(column));
} else if (VARCHAR.equals(type)) {
part.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);
part.bind(column, new Date(localMillis));
} else {
throw new IllegalArgumentException("Unsupported type " + type);
}
}
}
batch.execute();
}
}
private static void buildColumnInfo(
ConnectorTableMetadata tableMetadata,
ImmutableList.Builder<String> names,
ImmutableList.Builder<Type> types) {
for (ColumnMetadata column : tableMetadata.getColumns()) {
// TODO: also verify that the OutputFormat supports the type
if (!HiveRecordSink.isTypeSupported(column.getType())) {
throw new PrestoException(
NOT_SUPPORTED,
format(
"Cannot create table with unsupported type: %s",
column.getType().getDisplayName()));
}
names.add(column.getName());
types.add(column.getType());
}
if (tableMetadata.isSampled()) {
names.add(SAMPLE_WEIGHT_COLUMN_NAME);
types.add(BIGINT);
}
}
@Override
public List<Type> getColumnTypes() {
return ImmutableList.copyOf(transform(TASK_TABLE.getColumns(), ColumnMetadata::getType));
}
@Test
public void testGetRecords() throws Exception {
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);
ConnectorPartitionResult partitionResult =
splitManager.getPartitions(SESSION, tableHandle, TupleDomain.<ColumnHandle>all());
List<ConnectorSplit> splits =
getAllSplits(
splitManager.getPartitionSplits(SESSION, tableHandle, partitionResult.getPartitions()));
long rowNumber = 0;
for (ConnectorSplit split : splits) {
CassandraSplit cassandraSplit = (CassandraSplit) split;
long completedBytes = 0;
try (RecordCursor cursor =
recordSetProvider.getRecordSet(SESSION, cassandraSplit, 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, String.format("key %d", rowId));
assertEquals(
Bytes.toHexString(cursor.getSlice(columnIndex.get("typebytes")).getBytes()),
String.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(
cursor.getSlice(columnIndex.get("typeuuid")).toStringUtf8(),
String.format("00000000-0000-0000-0000-%012d", rowId));
assertEquals(
cursor.getSlice(columnIndex.get("typetimestamp")).toStringUtf8(),
Long.valueOf(DATE.getTime()).toString());
long newCompletedBytes = cursor.getCompletedBytes();
assertTrue(newCompletedBytes >= completedBytes);
completedBytes = newCompletedBytes;
}
}
}
assertEquals(rowNumber, 9);
}
private void doCreateSampledTable() throws InterruptedException {
// begin creating the table
List<ColumnMetadata> columns =
ImmutableList.<ColumnMetadata>builder()
.add(new ColumnMetadata("sales", BIGINT, 1, false))
.build();
ConnectorTableMetadata tableMetadata =
new ConnectorTableMetadata(temporaryCreateSampledTable, columns, tableOwner, true);
ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(SESSION, tableMetadata);
// write the records
RecordSink sink = recordSinkProvider.getRecordSink(outputHandle);
sink.beginRecord(8);
sink.appendLong(2);
sink.finishRecord();
sink.beginRecord(5);
sink.appendLong(3);
sink.finishRecord();
sink.beginRecord(7);
sink.appendLong(4);
sink.finishRecord();
String fragment = sink.commit();
// commit the table
metadata.commitCreateTable(outputHandle, ImmutableList.of(fragment));
// load the new table
ConnectorTableHandle tableHandle = getTableHandle(temporaryCreateSampledTable);
List<ConnectorColumnHandle> columnHandles =
ImmutableList.<ConnectorColumnHandle>builder()
.addAll(metadata.getColumnHandles(tableHandle).values())
.add(metadata.getSampleWeightColumnHandle(tableHandle))
.build();
assertEquals(columnHandles.size(), 2);
// verify the metadata
tableMetadata = metadata.getTableMetadata(getTableHandle(temporaryCreateSampledTable));
assertEquals(tableMetadata.getOwner(), tableOwner);
Map<String, ColumnMetadata> columnMap =
uniqueIndex(tableMetadata.getColumns(), columnNameGetter());
assertEquals(columnMap.size(), 1);
assertPrimitiveField(columnMap, 0, "sales", BIGINT, false);
// verify the data
ConnectorPartitionResult partitionResult =
splitManager.getPartitions(tableHandle, TupleDomain.<ConnectorColumnHandle>all());
assertEquals(partitionResult.getPartitions().size(), 1);
ConnectorSplitSource splitSource =
splitManager.getPartitionSplits(tableHandle, partitionResult.getPartitions());
ConnectorSplit split = getOnlyElement(getAllSplits(splitSource));
try (RecordCursor cursor = recordSetProvider.getRecordSet(split, columnHandles).cursor()) {
assertRecordCursorType(cursor, "rcfile-binary");
assertTrue(cursor.advanceNextPosition());
assertEquals(cursor.getLong(0), 2);
assertEquals(cursor.getLong(1), 8);
assertTrue(cursor.advanceNextPosition());
assertEquals(cursor.getLong(0), 3);
assertEquals(cursor.getLong(1), 5);
assertTrue(cursor.advanceNextPosition());
assertEquals(cursor.getLong(0), 4);
assertEquals(cursor.getLong(1), 7);
assertFalse(cursor.advanceNextPosition());
}
}
@Test
public void testGetRecords() throws Exception {
ConnectorTableHandle tableHandle = getTableHandle(table);
ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(tableHandle);
List<ConnectorColumnHandle> columnHandles =
ImmutableList.copyOf(metadata.getColumnHandles(tableHandle).values());
Map<String, Integer> columnIndex = indexColumns(columnHandles);
ConnectorPartitionResult partitionResult =
splitManager.getPartitions(tableHandle, TupleDomain.<ConnectorColumnHandle>all());
List<ConnectorSplit> splits =
getAllSplits(splitManager.getPartitionSplits(tableHandle, partitionResult.getPartitions()));
assertEquals(splits.size(), this.partitions.size());
for (ConnectorSplit split : splits) {
HiveSplit hiveSplit = (HiveSplit) split;
List<HivePartitionKey> partitionKeys = hiveSplit.getPartitionKeys();
String ds = partitionKeys.get(0).getValue();
String fileType = partitionKeys.get(1).getValue();
long dummy = Long.parseLong(partitionKeys.get(2).getValue());
long baseValue = getBaseValueForFileType(fileType);
assertEquals(dummy * 100, baseValue);
long rowNumber = 0;
long completedBytes = 0;
try (RecordCursor cursor =
recordSetProvider.getRecordSet(hiveSplit, columnHandles).cursor()) {
assertRecordCursorType(cursor, fileType);
assertEquals(cursor.getTotalBytes(), hiveSplit.getLength());
while (cursor.advanceNextPosition()) {
try {
assertReadFields(cursor, tableMetadata.getColumns());
} catch (RuntimeException e) {
throw new RuntimeException("row " + rowNumber, e);
}
rowNumber++;
if (rowNumber % 19 == 0) {
assertTrue(cursor.isNull(columnIndex.get("t_string")));
} else if (rowNumber % 19 == 1) {
assertEquals(cursor.getSlice(columnIndex.get("t_string")).toStringUtf8(), "");
} else {
assertEquals(
cursor.getSlice(columnIndex.get("t_string")).toStringUtf8(), (fileType + " test"));
}
assertEquals(
cursor.getLong(columnIndex.get("t_tinyint")),
(long) ((byte) (baseValue + 1 + rowNumber)));
assertEquals(cursor.getLong(columnIndex.get("t_smallint")), baseValue + 2 + rowNumber);
assertEquals(cursor.getLong(columnIndex.get("t_int")), baseValue + 3 + rowNumber);
if (rowNumber % 13 == 0) {
assertTrue(cursor.isNull(columnIndex.get("t_bigint")));
} else {
assertEquals(cursor.getLong(columnIndex.get("t_bigint")), baseValue + 4 + rowNumber);
}
assertEquals(
cursor.getDouble(columnIndex.get("t_float")), baseValue + 5.1 + rowNumber, 0.001);
assertEquals(cursor.getDouble(columnIndex.get("t_double")), baseValue + 6.2 + rowNumber);
if (rowNumber % 3 == 2) {
assertTrue(cursor.isNull(columnIndex.get("t_boolean")));
} else {
assertEquals(cursor.getBoolean(columnIndex.get("t_boolean")), rowNumber % 3 != 0);
}
if (rowNumber % 17 == 0) {
assertTrue(cursor.isNull(columnIndex.get("t_timestamp")));
} else {
long millis = new DateTime(2011, 5, 6, 7, 8, 9, 123, timeZone).getMillis();
assertEquals(
cursor.getLong(columnIndex.get("t_timestamp")), millis, (fileType + " test"));
}
if (rowNumber % 23 == 0) {
assertTrue(cursor.isNull(columnIndex.get("t_binary")));
} else {
assertEquals(
cursor.getSlice(columnIndex.get("t_binary")).toStringUtf8(), (fileType + " test"));
}
if (rowNumber % 29 == 0) {
assertTrue(cursor.isNull(columnIndex.get("t_map")));
} else {
String expectedJson = "{\"format\":\"" + fileType + "\"}";
assertEquals(cursor.getSlice(columnIndex.get("t_map")).toStringUtf8(), expectedJson);
}
if (rowNumber % 27 == 0) {
assertTrue(cursor.isNull(columnIndex.get("t_array_string")));
} else {
String expectedJson = "[\"" + fileType + "\",\"test\",\"data\"]";
assertEquals(
cursor.getSlice(columnIndex.get("t_array_string")).toStringUtf8(), expectedJson);
}
if (rowNumber % 31 == 0) {
assertTrue(cursor.isNull(columnIndex.get("t_complex")));
} else {
String expectedJson =
"{\"1\":[{\"s_string\":\""
+ fileType
+ "-a\",\"s_double\":0.1},{\"s_string\":\""
+ fileType
+ "-b\",\"s_double\":0.2}]}";
assertEquals(
cursor.getSlice(columnIndex.get("t_complex")).toStringUtf8(), expectedJson);
}
assertEquals(cursor.getSlice(columnIndex.get("ds")).toStringUtf8(), ds);
assertEquals(cursor.getSlice(columnIndex.get("file_format")).toStringUtf8(), fileType);
assertEquals(cursor.getLong(columnIndex.get("dummy")), dummy);
long newCompletedBytes = cursor.getCompletedBytes();
assertTrue(newCompletedBytes >= completedBytes);
assertTrue(newCompletedBytes <= hiveSplit.getLength());
completedBytes = newCompletedBytes;
}
}
assertTrue(completedBytes <= hiveSplit.getLength());
assertEquals(rowNumber, 100);
}
}
private void doCreateTable() throws InterruptedException {
// begin creating the table
List<ColumnMetadata> columns =
ImmutableList.<ColumnMetadata>builder()
.add(new ColumnMetadata("id", BIGINT, 1, false))
.add(new ColumnMetadata("t_string", VARCHAR, 2, false))
.add(new ColumnMetadata("t_bigint", BIGINT, 3, false))
.add(new ColumnMetadata("t_double", DOUBLE, 4, false))
.add(new ColumnMetadata("t_boolean", BOOLEAN, 5, false))
.build();
ConnectorTableMetadata tableMetadata =
new ConnectorTableMetadata(temporaryCreateTable, columns, tableOwner);
ConnectorOutputTableHandle outputHandle = metadata.beginCreateTable(SESSION, tableMetadata);
// write the records
RecordSink sink = recordSinkProvider.getRecordSink(outputHandle);
sink.beginRecord(1);
sink.appendLong(1);
sink.appendString("hello".getBytes(UTF_8));
sink.appendLong(123);
sink.appendDouble(43.5);
sink.appendBoolean(true);
sink.finishRecord();
sink.beginRecord(1);
sink.appendLong(2);
sink.appendNull();
sink.appendNull();
sink.appendNull();
sink.appendNull();
sink.finishRecord();
sink.beginRecord(1);
sink.appendLong(3);
sink.appendString("bye".getBytes(UTF_8));
sink.appendLong(456);
sink.appendDouble(98.1);
sink.appendBoolean(false);
sink.finishRecord();
String fragment = sink.commit();
// commit the table
metadata.commitCreateTable(outputHandle, ImmutableList.of(fragment));
// load the new table
ConnectorTableHandle tableHandle = getTableHandle(temporaryCreateTable);
List<ConnectorColumnHandle> columnHandles =
ImmutableList.copyOf(metadata.getColumnHandles(tableHandle).values());
// verify the metadata
tableMetadata = metadata.getTableMetadata(getTableHandle(temporaryCreateTable));
assertEquals(tableMetadata.getOwner(), tableOwner);
Map<String, ColumnMetadata> columnMap =
uniqueIndex(tableMetadata.getColumns(), columnNameGetter());
assertPrimitiveField(columnMap, 0, "id", BIGINT, false);
assertPrimitiveField(columnMap, 1, "t_string", VARCHAR, false);
assertPrimitiveField(columnMap, 2, "t_bigint", BIGINT, false);
assertPrimitiveField(columnMap, 3, "t_double", DOUBLE, false);
assertPrimitiveField(columnMap, 4, "t_boolean", BOOLEAN, false);
// verify the data
ConnectorPartitionResult partitionResult =
splitManager.getPartitions(tableHandle, TupleDomain.<ConnectorColumnHandle>all());
assertEquals(partitionResult.getPartitions().size(), 1);
ConnectorSplitSource splitSource =
splitManager.getPartitionSplits(tableHandle, partitionResult.getPartitions());
ConnectorSplit split = getOnlyElement(getAllSplits(splitSource));
try (RecordCursor cursor = recordSetProvider.getRecordSet(split, columnHandles).cursor()) {
assertRecordCursorType(cursor, "rcfile-binary");
assertTrue(cursor.advanceNextPosition());
assertEquals(cursor.getLong(0), 1);
assertEquals(cursor.getSlice(1).toStringUtf8(), "hello");
assertEquals(cursor.getLong(2), 123);
assertEquals(cursor.getDouble(3), 43.5);
assertEquals(cursor.getBoolean(4), true);
assertTrue(cursor.advanceNextPosition());
assertEquals(cursor.getLong(0), 2);
assertTrue(cursor.isNull(1));
assertTrue(cursor.isNull(2));
assertTrue(cursor.isNull(3));
assertTrue(cursor.isNull(4));
assertTrue(cursor.advanceNextPosition());
assertEquals(cursor.getLong(0), 3);
assertEquals(cursor.getSlice(1).toStringUtf8(), "bye");
assertEquals(cursor.getLong(2), 456);
assertEquals(cursor.getDouble(3), 98.1);
assertEquals(cursor.getBoolean(4), false);
assertFalse(cursor.advanceNextPosition());
}
}