public InMemoryRecordSet build() { return new InMemoryRecordSet(types, records); } }
@Override public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint) { TupleDomain<ColumnHandle> targetTupleDomain = constraint.transform(fieldIdToColumnHandle::get); Predicate<Map<ColumnHandle, NullableValue>> targetPredicate = convertToPredicate(targetTupleDomain); Constraint<ColumnHandle> targetConstraint = new Constraint<>(targetTupleDomain, targetPredicate); Iterable<List<Object>> records = () -> stream(partitionManager.getPartitions(metastore, sourceTableHandle, targetConstraint).getPartitions()) .map(hivePartition -> (List<Object>) IntStream.range(0, partitionColumns.size()) .mapToObj(fieldIdToColumnHandle::get) .map(columnHandle -> hivePartition.getKeys().get(columnHandle).getValue()) .collect(toList())) .iterator(); return new InMemoryRecordSet(partitionColumnTypes, records).cursor(); } };
@Override public RecordSet getRecordSet(ConnectorTransactionHandle transaction, ConnectorSession session, ConnectorSplit split, List<? extends ColumnHandle> columns) { JmxTableHandle tableHandle = ((JmxSplit) split).getTableHandle(); requireNonNull(columns, "columns is null"); checkArgument(!columns.isEmpty(), "must provide at least one column"); List<List<Object>> rows; try { if (tableHandle.isLiveData()) { rows = getLiveRows(tableHandle, columns); } else { List<Integer> selectedColumns = calculateSelectedColumns(tableHandle.getColumnHandles(), getColumnNames(columns)); rows = tableHandle.getObjectNames().stream() .flatMap(objectName -> jmxHistoricalData.getRows(objectName, selectedColumns).stream()) .collect(toImmutableList()); } } catch (JMException e) { rows = ImmutableList.of(); } return new InMemoryRecordSet(getColumnTypes(columns), rows); }
@Test public void test() { TypeManager typeManager = new TypeRegistry(); FunctionRegistry functionRegistry = new FunctionRegistry(typeManager, new BlockEncodingManager(typeManager), new FeaturesConfig()); ArrayType arrayOfBigintType = new ArrayType(BIGINT); FieldSetFilteringRecordSet fieldSetFilteringRecordSet = new FieldSetFilteringRecordSet( functionRegistry, new InMemoryRecordSet( ImmutableList.of(BIGINT, BIGINT, TIMESTAMP_WITH_TIME_ZONE, TIMESTAMP_WITH_TIME_ZONE, arrayOfBigintType, arrayOfBigintType), ImmutableList.of( ImmutableList.of( 100L, 100L, // test same time in different time zone to make sure equal check was done properly packDateTimeWithZone(100, getTimeZoneKeyForOffset(123)), packDateTimeWithZone(100, getTimeZoneKeyForOffset(234)), // test structural type arrayBlockOf(BIGINT, 12, 34, 56), arrayBlockOf(BIGINT, 12, 34, 56)))), ImmutableList.of(ImmutableSet.of(0, 1), ImmutableSet.of(2, 3), ImmutableSet.of(4, 5))); RecordCursor recordCursor = fieldSetFilteringRecordSet.cursor(); assertTrue(recordCursor.advanceNextPosition()); } }
ImmutableList.of(column("a", INTEGER)), ImmutableList.of(column("b", INTEGER)), new InMemoryRecordSet(ImmutableList.of(INTEGER), generateKeys(1, splits + 1)), MAX_BYTES_PER_RESPONSE, lookupRequestsConcurrency);
ImmutableList.of(column("a", INTEGER)), ImmutableList.of(column("b", INTEGER)), new InMemoryRecordSet(ImmutableList.of(INTEGER), generateKeys(0, splits)), MAX_BYTES_PER_RESPONSE, lookupRequestsConcurrency);
public InMemoryRecordSet build() { return new InMemoryRecordSet(types, records); } }
@Test public void test() { ArrayType arrayOfBigintType = new ArrayType(BIGINT); FieldSetFilteringRecordSet fieldSetFilteringRecordSet = new FieldSetFilteringRecordSet( FUNCTION_REGISTRY, new InMemoryRecordSet( ImmutableList.of(BIGINT, BIGINT, TIMESTAMP_WITH_TIME_ZONE, TIMESTAMP_WITH_TIME_ZONE, arrayOfBigintType, arrayOfBigintType), ImmutableList.of( ImmutableList.of( 100L, 100L, // test same time in different time zone to make sure equal check was done properly packDateTimeWithZone(100, getTimeZoneKeyForOffset(123)), packDateTimeWithZone(100, getTimeZoneKeyForOffset(234)), // test structural type arrayBlockOf(BIGINT, 12, 34, 56), arrayBlockOf(BIGINT, 12, 34, 56) ))), ImmutableList.of(ImmutableSet.of(0, 1), ImmutableSet.of(2, 3), ImmutableSet.of(4, 5))); RecordCursor recordCursor = fieldSetFilteringRecordSet.cursor(); assertTrue(recordCursor.advanceNextPosition()); } }
@Test public void testSingleColumn() throws Exception { InMemoryRecordSet records = new InMemoryRecordSet(ImmutableList.of(VARCHAR), ImmutableList.copyOf(new List<?>[] {ImmutableList.of("abc"), ImmutableList.of("def"), ImmutableList.of("g")})); OperatorContext operatorContext = driverContext.addOperatorContext(0, new PlanNodeId("test"), RecordProjectOperator.class.getSimpleName()); Operator operator = new RecordProjectOperator(operatorContext, records); MaterializedResult expected = resultBuilder(driverContext.getSession(), VARCHAR) .row("abc") .row("def") .row("g") .build(); OperatorAssertion.assertOperatorEquals(operator, expected); }
@Test public void testMultiColumn() throws Exception { InMemoryRecordSet records = new InMemoryRecordSet(ImmutableList.of(VARCHAR, BIGINT), ImmutableList.of( ImmutableList.of("abc", 1L), ImmutableList.of("def", 2L), ImmutableList.of("g", 0L))); OperatorContext operatorContext = driverContext.addOperatorContext(0, new PlanNodeId("test"), RecordProjectOperator.class.getSimpleName()); Operator operator = new RecordProjectOperator(operatorContext, records); MaterializedResult expected = resultBuilder(driverContext.getSession(), VARCHAR, BIGINT) .row("abc", 1) .row("def", 2) .row("g", 0) .build(); OperatorAssertion.assertOperatorEquals(operator, expected); }