use of io.trino.split.MappedPageSource in project trino by trinodb.
the class SystemPageSourceProvider method createPageSource.
@Override
public ConnectorPageSource createPageSource(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, ConnectorTableHandle table, List<ColumnHandle> columns, DynamicFilter dynamicFilter) {
requireNonNull(columns, "columns is null");
SystemTransactionHandle systemTransaction = (SystemTransactionHandle) transaction;
SystemSplit systemSplit = (SystemSplit) split;
SchemaTableName tableName = ((SystemTableHandle) table).getSchemaTableName();
SystemTable systemTable = tables.getSystemTable(session, tableName).orElseThrow(() -> new TrinoException(NOT_FOUND, format("Table '%s' not found", tableName)));
List<ColumnMetadata> tableColumns = systemTable.getTableMetadata().getColumns();
Map<String, Integer> columnsByName = new HashMap<>();
for (int i = 0; i < tableColumns.size(); i++) {
ColumnMetadata column = tableColumns.get(i);
if (columnsByName.put(column.getName(), i) != null) {
throw new TrinoException(GENERIC_INTERNAL_ERROR, "Duplicate column name: " + column.getName());
}
}
ImmutableList.Builder<Integer> userToSystemFieldIndex = ImmutableList.builder();
for (ColumnHandle column : columns) {
String columnName = ((SystemColumnHandle) column).getColumnName();
Integer index = columnsByName.get(columnName);
if (index == null) {
throw new TrinoException(GENERIC_INTERNAL_ERROR, format("Column does not exist: %s.%s", tableName, columnName));
}
userToSystemFieldIndex.add(index);
}
TupleDomain<ColumnHandle> constraint = systemSplit.getConstraint();
if (constraint.isNone()) {
return new EmptyPageSource();
}
TupleDomain<Integer> newConstraint = systemSplit.getConstraint().transformKeys(columnHandle -> columnsByName.get(((SystemColumnHandle) columnHandle).getColumnName()));
try {
return new MappedPageSource(systemTable.pageSource(systemTransaction.getConnectorTransactionHandle(), session, newConstraint), userToSystemFieldIndex.build());
} catch (UnsupportedOperationException e) {
return new RecordPageSource(new MappedRecordSet(toRecordSet(systemTransaction.getConnectorTransactionHandle(), systemTable, session, newConstraint), userToSystemFieldIndex.build()));
}
}
Aggregations