Search in sources :

Example 1 with RecordSet

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

the class IndexSourceOperator method addSplit.

@Override
public Supplier<Optional<UpdatablePageSource>> addSplit(Split split) {
    requireNonNull(split, "split is null");
    checkState(source == null, "Index source split already set");
    IndexSplit indexSplit = (IndexSplit) split.getConnectorSplit();
    // Normalize the incoming RecordSet to something that can be consumed by the index
    RecordSet normalizedRecordSet = probeKeyNormalizer.apply(indexSplit.getKeyRecordSet());
    ConnectorPageSource result = index.lookup(normalizedRecordSet);
    source = new PageSourceOperator(result, operatorContext);
    Object splitInfo = split.getInfo();
    if (splitInfo != null) {
        operatorContext.setInfoSupplier(Suppliers.ofInstance(new SplitOperatorInfo(split.getCatalogName(), splitInfo)));
    }
    return Optional::empty;
}
Also used : SplitOperatorInfo(io.trino.operator.SplitOperatorInfo) RecordSet(io.trino.spi.connector.RecordSet) ConnectorPageSource(io.trino.spi.connector.ConnectorPageSource) PageSourceOperator(io.trino.operator.PageSourceOperator)

Example 2 with RecordSet

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

the class TestExampleRecordSet method testGetColumnTypes.

@Test
public void testGetColumnTypes() {
    RecordSet recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(new ExampleColumnHandle("text", createUnboundedVarcharType(), 0), new ExampleColumnHandle("value", BIGINT, 1)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(createUnboundedVarcharType(), BIGINT));
    recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(new ExampleColumnHandle("value", BIGINT, 1), new ExampleColumnHandle("text", createUnboundedVarcharType(), 0)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, createUnboundedVarcharType()));
    recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(new ExampleColumnHandle("value", BIGINT, 1), new ExampleColumnHandle("value", BIGINT, 1), new ExampleColumnHandle("text", createUnboundedVarcharType(), 0)));
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of(BIGINT, BIGINT, createUnboundedVarcharType()));
    recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of());
    assertEquals(recordSet.getColumnTypes(), ImmutableList.of());
}
Also used : RecordSet(io.trino.spi.connector.RecordSet) Test(org.testng.annotations.Test)

Example 3 with RecordSet

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

the class TestExampleRecordSet method testCursorSimple.

@Test
public void testCursorSimple() {
    RecordSet recordSet = new ExampleRecordSet(new ExampleSplit(dataUri), ImmutableList.of(new ExampleColumnHandle("text", createUnboundedVarcharType(), 0), new ExampleColumnHandle("value", BIGINT, 1)));
    RecordCursor cursor = recordSet.cursor();
    assertEquals(cursor.getType(0), createUnboundedVarcharType());
    assertEquals(cursor.getType(1), BIGINT);
    Map<String, Long> data = new LinkedHashMap<>();
    while (cursor.advanceNextPosition()) {
        data.put(cursor.getSlice(0).toStringUtf8(), cursor.getLong(1));
        assertFalse(cursor.isNull(0));
        assertFalse(cursor.isNull(1));
    }
    assertEquals(data, ImmutableMap.<String, Long>builder().put("ten", 10L).put("eleven", 11L).put("twelve", 12L).buildOrThrow());
}
Also used : RecordCursor(io.trino.spi.connector.RecordCursor) RecordSet(io.trino.spi.connector.RecordSet) LinkedHashMap(java.util.LinkedHashMap) Test(org.testng.annotations.Test)

Example 4 with RecordSet

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

the class TestJdbcRecordSet method testIdempotentClose.

@Test
public void testIdempotentClose() {
    RecordSet recordSet = createRecordSet(ImmutableList.of(columnHandles.get("value"), columnHandles.get("value"), columnHandles.get("text")));
    RecordCursor cursor = recordSet.cursor();
    cursor.close();
    cursor.close();
}
Also used : RecordCursor(io.trino.spi.connector.RecordCursor) RecordSet(io.trino.spi.connector.RecordSet) Test(org.testng.annotations.Test)

Example 5 with RecordSet

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

the class TestJdbcRecordSetProvider method testGetRecordSet.

@Test
public void testGetRecordSet() {
    ConnectorTransactionHandle transaction = new JdbcTransactionHandle();
    JdbcRecordSetProvider recordSetProvider = new JdbcRecordSetProvider(jdbcClient, executor);
    RecordSet recordSet = recordSetProvider.getRecordSet(transaction, SESSION, split, table, ImmutableList.of(textColumn, textShortColumn, valueColumn));
    assertNotNull(recordSet, "recordSet is null");
    RecordCursor cursor = recordSet.cursor();
    assertNotNull(cursor, "cursor is null");
    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));
    }
    assertEquals(data, ImmutableMap.<String, Long>builder().put("one", 1L).put("two", 2L).put("three", 3L).put("ten", 10L).put("eleven", 11L).put("twelve", 12L).buildOrThrow());
}
Also used : RecordCursor(io.trino.spi.connector.RecordCursor) ConnectorTransactionHandle(io.trino.spi.connector.ConnectorTransactionHandle) OptionalLong(java.util.OptionalLong) RecordSet(io.trino.spi.connector.RecordSet) LinkedHashMap(java.util.LinkedHashMap) Test(org.testng.annotations.Test)

Aggregations

RecordSet (io.trino.spi.connector.RecordSet)25 Test (org.testng.annotations.Test)14 RecordCursor (io.trino.spi.connector.RecordCursor)12 LinkedHashMap (java.util.LinkedHashMap)8 ConnectorTransactionHandle (io.trino.spi.connector.ConnectorTransactionHandle)6 List (java.util.List)6 ImmutableList.toImmutableList (com.google.common.collect.ImmutableList.toImmutableList)5 ColumnHandle (io.trino.spi.connector.ColumnHandle)5 Objects.requireNonNull (java.util.Objects.requireNonNull)5 ImmutableSet.toImmutableSet (com.google.common.collect.ImmutableSet.toImmutableSet)4 ConnectorSession (io.trino.spi.connector.ConnectorSession)4 ConnectorTableHandle (io.trino.spi.connector.ConnectorTableHandle)4 ArrayList (java.util.ArrayList)4 Preconditions.checkArgument (com.google.common.base.Preconditions.checkArgument)3 Preconditions.checkState (com.google.common.base.Preconditions.checkState)3 ImmutableList (com.google.common.collect.ImmutableList)3 ImmutableMap (com.google.common.collect.ImmutableMap)3 DispatchingRowDecoderFactory (io.trino.decoder.DispatchingRowDecoderFactory)3 RowDecoder (io.trino.decoder.RowDecoder)3 ConnectorRecordSetProvider (io.trino.spi.connector.ConnectorRecordSetProvider)3