Examples with ConnectorSession com.facebook.presto.spi.ConnectorSession used on opensource projects

Example 36 with ConnectorSession

use of com.facebook.presto.spi.ConnectorSession in project presto by prestodb.

the class TpchIndexProvider method getIndex.

@Override
public ConnectorIndex getIndex(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorIndexHandle indexHandle, List<ColumnHandle> lookupSchema, List<ColumnHandle> outputSchema) {
    TpchIndexHandle tpchIndexHandle = (TpchIndexHandle) indexHandle;
    Map<ColumnHandle, NullableValue> fixedValues = TupleDomain.extractFixedValues(tpchIndexHandle.getFixedValues()).get();
    checkArgument(lookupSchema.stream().noneMatch(handle -> fixedValues.keySet().contains(handle)), "Lookup columnHandles are not expected to overlap with the fixed value predicates");
    // Establish an order for the fixedValues
    List<ColumnHandle> fixedValueColumns = ImmutableList.copyOf(fixedValues.keySet());
    // Extract the fixedValues as their raw values and types
    List<Object> rawFixedValues = new ArrayList<>(fixedValueColumns.size());
    List<Type> rawFixedTypes = new ArrayList<>(fixedValueColumns.size());
    for (ColumnHandle fixedValueColumn : fixedValueColumns) {
        rawFixedValues.add(fixedValues.get(fixedValueColumn).getValue());
        rawFixedTypes.add(((TpchColumnHandle) fixedValueColumn).getType());
    }
    // Establish the schema after we append the fixed values to the lookup keys.
    List<ColumnHandle> finalLookupSchema = ImmutableList.<ColumnHandle>builder().addAll(lookupSchema).addAll(fixedValueColumns).build();
    Optional<TpchIndexedData.IndexedTable> indexedTable = indexedData.getIndexedTable(tpchIndexHandle.getTableName(), tpchIndexHandle.getScaleFactor(), tpchIndexHandle.getIndexColumnNames());
    checkState(indexedTable.isPresent());
    TpchIndexedData.IndexedTable table = indexedTable.get();
    // Compute how to map from the final lookup schema to the table index key order
    List<Integer> keyRemap = computeRemap(handleToNames(finalLookupSchema), table.getKeyColumns());
    Function<RecordSet, RecordSet> keyFormatter = key -> new MappedRecordSet(new AppendingRecordSet(key, rawFixedValues, rawFixedTypes), keyRemap);
    // Compute how to map from the output of the indexed data to the expected output schema
    List<Integer> outputRemap = computeRemap(table.getOutputColumns(), handleToNames(outputSchema));
    Function<RecordSet, RecordSet> outputFormatter = output -> new MappedRecordSet(output, outputRemap);
    return new TpchConnectorIndex(keyFormatter, outputFormatter, table);
}

Example 37 with ConnectorSession

use of com.facebook.presto.spi.ConnectorSession in project presto by prestodb.

the class TpchIndexMetadata method resolveIndex.

@Override
public Optional<ConnectorResolvedIndex> resolveIndex(ConnectorSession session, ConnectorTableHandle tableHandle, Set<ColumnHandle> indexableColumns, Set<ColumnHandle> outputColumns, TupleDomain<ColumnHandle> tupleDomain) {
    TpchTableHandle tpchTableHandle = (TpchTableHandle) tableHandle;
    // Keep the fixed values that don't overlap with the indexableColumns
    // Note: technically we could more efficiently utilize the overlapped columns, but this way is simpler for now
    Map<ColumnHandle, NullableValue> fixedValues = TupleDomain.extractFixedValues(tupleDomain).orElse(ImmutableMap.of()).entrySet().stream().filter(entry -> !indexableColumns.contains(entry.getKey())).filter(// strip nulls since meaningless in index join lookups
    entry -> !entry.getValue().isNull()).collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
    // determine all columns available for index lookup
    Set<String> lookupColumnNames = ImmutableSet.<String>builder().addAll(handleToNames(ImmutableList.copyOf(indexableColumns))).addAll(handleToNames(ImmutableList.copyOf(fixedValues.keySet()))).build();
    // do we have an index?
    if (!indexedData.getIndexedTable(tpchTableHandle.getTableName(), tpchTableHandle.getScaleFactor(), lookupColumnNames).isPresent()) {
        return Optional.empty();
    }
    TupleDomain<ColumnHandle> filteredTupleDomain = tupleDomain;
    if (!tupleDomain.isNone()) {
        filteredTupleDomain = TupleDomain.withColumnDomains(Maps.filterKeys(tupleDomain.getDomains().get(), not(in(fixedValues.keySet()))));
    }
    TpchIndexHandle indexHandle = new TpchIndexHandle(tpchTableHandle.getTableName(), tpchTableHandle.getScaleFactor(), lookupColumnNames, TupleDomain.fromFixedValues(fixedValues));
    return Optional.of(new ConnectorResolvedIndex(indexHandle, filteredTupleDomain));
}

Example 38 with ConnectorSession

use of com.facebook.presto.spi.ConnectorSession in project presto by prestodb.

the class ConnectorMetadata method getInsertLayout.

/**
 * Get the physical layout for a inserting into an existing table.
 */
default Optional<ConnectorNewTableLayout> getInsertLayout(ConnectorSession session, ConnectorTableHandle tableHandle) {
    List<ConnectorTableLayout> layouts = getTableLayouts(session, tableHandle, new Constraint<>(TupleDomain.all(), map -> true), Optional.empty()).stream().map(ConnectorTableLayoutResult::getTableLayout).filter(layout -> layout.getTablePartitioning().isPresent()).collect(toList());
    if (layouts.isEmpty()) {
        return Optional.empty();
    }
    if (layouts.size() > 1) {
        throw new PrestoException(NOT_SUPPORTED, "Tables with multiple layouts can not be written");
    }
    ConnectorTableLayout layout = layouts.get(0);
    ConnectorPartitioningHandle partitioningHandle = layout.getTablePartitioning().get().getPartitioningHandle();
    Map<ColumnHandle, String> columnNamesByHandle = getColumnHandles(session, tableHandle).entrySet().stream().collect(Collectors.toMap(Map.Entry::getValue, Map.Entry::getKey));
    List<String> partitionColumns = layout.getTablePartitioning().get().getPartitioningColumns().stream().map(columnNamesByHandle::get).collect(toList());
    return Optional.of(new ConnectorNewTableLayout(partitioningHandle, partitionColumns));
}

Example 39 with ConnectorSession

use of com.facebook.presto.spi.ConnectorSession in project presto by prestodb.

the class AbstractTestHiveClient method doTestTransactionDeleteInsert.

private void doTestTransactionDeleteInsert(HiveStorageFormat storageFormat, SchemaTableName tableName, Domain domainToDrop, MaterializedResult insertData, MaterializedResult expectedData, TransactionDeleteInsertTestTag tag, boolean expectQuerySucceed, Optional<ConflictTrigger> conflictTrigger) throws Exception {
    Path writePath = null;
    Path targetPath = null;
    try (Transaction transaction = newTransaction()) {
        try {
            ConnectorMetadata metadata = transaction.getMetadata();
            ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName);
            ConnectorSession session;
            rollbackIfEquals(tag, ROLLBACK_RIGHT_AWAY);
            // Query 1: delete
            session = newSession();
            HiveColumnHandle dsColumnHandle = (HiveColumnHandle) metadata.getColumnHandles(session, tableHandle).get("pk2");
            TupleDomain<ColumnHandle> tupleDomain = withColumnDomains(ImmutableMap.of(dsColumnHandle, domainToDrop));
            Constraint<ColumnHandle> constraint = new Constraint<>(tupleDomain, convertToPredicate(tupleDomain));
            ConnectorTableLayoutHandle tableLayoutHandle = getTableLayout(session, metadata, tableHandle, constraint, transaction).getHandle();
            metadata.metadataDelete(session, tableHandle, tableLayoutHandle);
            rollbackIfEquals(tag, ROLLBACK_AFTER_DELETE);
            // Query 2: insert
            session = newSession();
            ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle);
            rollbackIfEquals(tag, ROLLBACK_AFTER_BEGIN_INSERT);
            writePath = getStagingPathRoot(insertTableHandle);
            targetPath = getTargetPathRoot(insertTableHandle);
            ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle, TEST_HIVE_PAGE_SINK_CONTEXT);
            sink.appendPage(insertData.toPage());
            rollbackIfEquals(tag, ROLLBACK_AFTER_APPEND_PAGE);
            Collection<Slice> fragments = getFutureValue(sink.finish());
            rollbackIfEquals(tag, ROLLBACK_AFTER_SINK_FINISH);
            metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of());
            rollbackIfEquals(tag, ROLLBACK_AFTER_FINISH_INSERT);
            assertEquals(tag, COMMIT);
            if (conflictTrigger.isPresent()) {
                List<PartitionUpdate> partitionUpdates = fragments.stream().map(Slice::getBytes).map(HiveTestUtils.PARTITION_UPDATE_CODEC::fromJson).collect(toList());
                conflictTrigger.get().triggerConflict(session, tableName, insertTableHandle, partitionUpdates);
            }
            transaction.commit();
            if (conflictTrigger.isPresent()) {
                assertTrue(expectQuerySucceed);
                conflictTrigger.get().verifyAndCleanup(tableName);
            }
        } catch (TestingRollbackException e) {
            transaction.rollback();
        } catch (PrestoException e) {
            assertFalse(expectQuerySucceed);
            if (conflictTrigger.isPresent()) {
                conflictTrigger.get().verifyAndCleanup(tableName);
            }
        }
    }
    // check that temporary files are removed
    if (writePath != null && !writePath.equals(targetPath)) {
        HdfsContext context = new HdfsContext(newSession(), tableName.getSchemaName(), tableName.getTableName(), writePath.toString(), false);
        FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, writePath);
        assertFalse(fileSystem.exists(writePath));
    }
    try (Transaction transaction = newTransaction()) {
        ConnectorSession session = newSession();
        MetastoreContext metastoreContext = new MetastoreContext(session.getIdentity(), session.getQueryId(), session.getClientInfo(), session.getSource(), getMetastoreHeaders(session), false, DEFAULT_COLUMN_CONVERTER_PROVIDER);
        // verify partitions
        List<String> partitionNames = transaction.getMetastore().getPartitionNames(metastoreContext, tableName.getSchemaName(), tableName.getTableName()).orElseThrow(() -> new AssertionError("Table does not exist: " + tableName));
        assertEqualsIgnoreOrder(partitionNames, expectedData.getMaterializedRows().stream().map(row -> format("pk1=%s/pk2=%s", row.getField(1), row.getField(2))).distinct().collect(toList()));
        // load the new table
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName);
        List<ColumnHandle> columnHandles = filterNonHiddenColumnHandles(metadata.getColumnHandles(session, tableHandle).values());
        // verify the data
        MaterializedResult result = readTable(transaction, tableHandle, columnHandles, session, TupleDomain.all(), OptionalInt.empty(), Optional.of(storageFormat));
        assertEqualsIgnoreOrder(result.getMaterializedRows(), expectedData.getMaterializedRows());
    }
}

Example 40 with ConnectorSession

use of com.facebook.presto.spi.ConnectorSession in project presto by prestodb.

the class AbstractTestHiveClient method assertGetRecords.

protected void assertGetRecords(String tableName, HiveStorageFormat hiveStorageFormat) throws Exception {
    try (Transaction transaction = newTransaction()) {
        ConnectorSession session = newSession();
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorTableHandle tableHandle = getTableHandle(metadata, new SchemaTableName(database, tableName));
        ConnectorTableMetadata tableMetadata = metadata.getTableMetadata(session, tableHandle);
        ConnectorTableLayoutHandle layoutHandle = getLayout(session, transaction, tableHandle, TupleDomain.all());
        List<ConnectorSplit> splits = getAllSplits(session, transaction, layoutHandle);
        assertEquals(splits.size(), 1);
        HiveSplit hiveSplit = (HiveSplit) getOnlyElement(splits);
        List<ColumnHandle> columnHandles = ImmutableList.copyOf(metadata.getColumnHandles(session, tableHandle).values());
        ConnectorPageSource pageSource = pageSourceProvider.createPageSource(transaction.getTransactionHandle(), session, hiveSplit, layoutHandle, columnHandles, NON_CACHEABLE);
        assertGetRecords(hiveStorageFormat, tableMetadata, hiveSplit, pageSource, columnHandles);
    }
}