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

Example 16 with ConnectorPageSink

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

the class TestDeltaLakePageSink method testPageSinkStats.

@Test
public void testPageSinkStats() throws Exception {
    File tempDir = Files.createTempDir();
    try {
        DeltaLakeWriterStats stats = new DeltaLakeWriterStats();
        String tablePath = tempDir.getAbsolutePath() + "/test_table";
        ConnectorPageSink pageSink = createPageSink(new Path(tablePath), stats);
        List<LineItemColumn> columns = ImmutableList.copyOf(LineItemColumn.values());
        List<Type> columnTypes = columns.stream().map(LineItemColumn::getType).map(TestDeltaLakePageSink::getTrinoType).collect(toList());
        PageBuilder pageBuilder = new PageBuilder(columnTypes);
        long rows = 0;
        for (LineItem lineItem : new LineItemGenerator(0.01, 1, 1)) {
            if (rows >= NUM_ROWS) {
                break;
            }
            rows++;
            pageBuilder.declarePosition();
            for (int i = 0; i < columns.size(); i++) {
                LineItemColumn column = columns.get(i);
                BlockBuilder blockBuilder = pageBuilder.getBlockBuilder(i);
                writeToBlock(blockBuilder, column, lineItem);
            }
        }
        Page page = pageBuilder.build();
        pageSink.appendPage(page);
        JsonCodec<DataFileInfo> dataFileInfoCodec = new JsonCodecFactory().jsonCodec(DataFileInfo.class);
        Collection<Slice> fragments = getFutureValue(pageSink.finish());
        List<DataFileInfo> dataFileInfos = fragments.stream().map(Slice::getBytes).map(dataFileInfoCodec::fromJson).collect(toImmutableList());
        assertEquals(dataFileInfos.size(), 1);
        DataFileInfo dataFileInfo = dataFileInfos.get(0);
        List<File> files = ImmutableList.copyOf(new File(tablePath).listFiles((dir, name) -> !name.endsWith(".crc")));
        assertEquals(files.size(), 1);
        File outputFile = files.get(0);
        assertEquals(round(stats.getInputPageSizeInBytes().getAllTime().getMax()), page.getRetainedSizeInBytes());
        assertEquals(dataFileInfo.getStatistics().getNumRecords(), Optional.of(rows));
        assertEquals(dataFileInfo.getPartitionValues(), ImmutableList.of());
        assertEquals(dataFileInfo.getSize(), outputFile.length());
        assertEquals(dataFileInfo.getPath(), outputFile.getName());
        Instant now = Instant.now();
        assertTrue(dataFileInfo.getCreationTime() < now.toEpochMilli());
        assertTrue(dataFileInfo.getCreationTime() > now.minus(1, MINUTES).toEpochMilli());
    } finally {
        deleteRecursively(tempDir.toPath(), ALLOW_INSECURE);
    }
}

Example 17 with ConnectorPageSink

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

the class AbstractTestHive method insertData.

private String insertData(SchemaTableName tableName, MaterializedResult data, Map<String, Object> sessionProperties) throws Exception {
    Path writePath;
    Path targetPath;
    String queryId;
    try (Transaction transaction = newTransaction()) {
        ConnectorMetadata metadata = transaction.getMetadata();
        ConnectorSession session = newSession(sessionProperties);
        ConnectorTableHandle tableHandle = getTableHandle(metadata, tableName);
        ConnectorInsertTableHandle insertTableHandle = metadata.beginInsert(session, tableHandle, ImmutableList.of(), NO_RETRIES);
        queryId = session.getQueryId();
        writePath = getStagingPathRoot(insertTableHandle);
        targetPath = getTargetPathRoot(insertTableHandle);
        ConnectorPageSink sink = pageSinkProvider.createPageSink(transaction.getTransactionHandle(), session, insertTableHandle);
        // write data
        sink.appendPage(data.toPage());
        Collection<Slice> fragments = getFutureValue(sink.finish());
        // commit the insert
        metadata.finishInsert(session, insertTableHandle, fragments, ImmutableList.of());
        transaction.commit();
    }
    // check that temporary files are removed
    if (!writePath.equals(targetPath)) {
        HdfsContext context = new HdfsContext(newSession());
        FileSystem fileSystem = hdfsEnvironment.getFileSystem(context, writePath);
        assertFalse(fileSystem.exists(writePath));
    }
    return queryId;
}

Example 18 with ConnectorPageSink

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

the class TestMemoryPagesStore method createTable.

private void createTable(long tableId, Long... activeTableIds) {
    ConnectorPageSink pageSink = pageSinkProvider.createPageSink(MemoryTransactionHandle.INSTANCE, SESSION, createMemoryOutputTableHandle(tableId, activeTableIds));
    pageSink.finish();
}