Examples with DataContext org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext used on opensource projects

Example 31 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project calcite by apache.

the class ScannableTableTest method testPrepared2.

/**
 * Test case for
 * <a href="https://issues.apache.org/jira/browse/CALCITE-1031">[CALCITE-1031]
 * In prepared statement, CsvScannableTable.scan is called twice</a>.
 */
@Test
public void testPrepared2() throws SQLException {
    final Properties properties = new Properties();
    properties.setProperty("caseSensitive", "true");
    try (final Connection connection = DriverManager.getConnection("jdbc:calcite:", properties)) {
        final CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);
        final AtomicInteger scanCount = new AtomicInteger();
        final AtomicInteger enumerateCount = new AtomicInteger();
        final Schema schema = new AbstractSchema() {

            @Override
            protected Map<String, Table> getTableMap() {
                return ImmutableMap.<String, Table>of("TENS", new SimpleTable() {

                    private Enumerable<Object[]> superScan(DataContext root) {
                        return super.scan(root);
                    }

                    @Override
                    public Enumerable<Object[]> scan(final DataContext root) {
                        scanCount.incrementAndGet();
                        return new AbstractEnumerable<Object[]>() {

                            public Enumerator<Object[]> enumerator() {
                                enumerateCount.incrementAndGet();
                                return superScan(root).enumerator();
                            }
                        };
                    }
                });
            }
        };
        calciteConnection.getRootSchema().add("TEST", schema);
        final String sql = "select * from \"TEST\".\"TENS\" where \"i\" < ?";
        final PreparedStatement statement = calciteConnection.prepareStatement(sql);
        assertThat(scanCount.get(), is(0));
        assertThat(enumerateCount.get(), is(0));
        // First execute
        statement.setInt(1, 20);
        assertThat(scanCount.get(), is(0));
        ResultSet resultSet = statement.executeQuery();
        assertThat(scanCount.get(), is(1));
        assertThat(enumerateCount.get(), is(1));
        assertThat(resultSet, Matchers.returnsUnordered("i=0", "i=10"));
        assertThat(scanCount.get(), is(1));
        assertThat(enumerateCount.get(), is(1));
        // Second execute
        resultSet = statement.executeQuery();
        assertThat(scanCount.get(), is(2));
        assertThat(resultSet, Matchers.returnsUnordered("i=0", "i=10"));
        assertThat(scanCount.get(), is(2));
        // Third execute
        statement.setInt(1, 30);
        resultSet = statement.executeQuery();
        assertThat(scanCount.get(), is(3));
        assertThat(resultSet, Matchers.returnsUnordered("i=0", "i=10", "i=20"));
        assertThat(scanCount.get(), is(3));
    }
}

Example 32 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project calcite by apache.

the class PsTableFunction method eval.

public static ScannableTable eval(boolean b) {
    return new ScannableTable() {

        public Enumerable<Object[]> scan(DataContext root) {
            final RelDataType rowType = getRowType(root.getTypeFactory());
            final List<String> fieldNames = ImmutableList.copyOf(rowType.getFieldNames());
            final String[] args;
            final String osName = System.getProperty("os.name");
            final String osVersion = System.getProperty("os.version");
            Util.discard(osVersion);
            switch(osName) {
                case // tested on version 10.12.5
                "Mac OS X":
                    args = new String[] { "ps", "ax", "-o", "ppid=,pid=,pgid=,tpgid=,stat=," + "user=,pcpu=,pmem=,vsz=,rss=,tty=,start=,time=,uid=,ruid=," + "sess=,comm=" };
                    break;
                default:
                    args = new String[] { "ps", "--no-headers", "axo", "ppid,pid,pgrp," + "tpgid,stat,user,pcpu,pmem,vsz,rss,tty,start_time,time,euid," + "ruid,sess,comm" };
            }
            return Processes.processLines(args).select(new Function1<String, Object[]>() {

                public Object[] apply(String line) {
                    final String[] fields = line.trim().split(" +");
                    final Object[] values = new Object[fieldNames.size()];
                    for (int i = 0; i < values.length; i++) {
                        try {
                            values[i] = field(fieldNames.get(i), fields[i]);
                        } catch (RuntimeException e) {
                            throw new RuntimeException("while parsing value [" + fields[i] + "] of field [" + fieldNames.get(i) + "] in line [" + line + "]");
                        }
                    }
                    return values;
                }

                private Object field(String field, String value) {
                    switch(field) {
                        case "pid":
                        case "ppid":
                        // linux only; macOS equivalent is "pgid"
                        case "pgrp":
                        // see "pgrp"
                        case "pgid":
                        case "tpgid":
                            return Integer.valueOf(value);
                        case "pcpu":
                        case "pmem":
                            return (int) (Float.valueOf(value) * 10f);
                        case "time":
                            final Matcher m1 = MINUTE_SECOND_MILLIS_PATTERN.matcher(value);
                            if (m1.matches()) {
                                final long h = Long.parseLong(m1.group(1));
                                final long m = Long.parseLong(m1.group(2));
                                final long s = Long.parseLong(m1.group(3));
                                return h * 3600000L + m * 60000L + s * 1000L;
                            }
                            final Matcher m2 = HOUR_MINUTE_SECOND_PATTERN.matcher(value);
                            if (m2.matches()) {
                                final long m = Long.parseLong(m2.group(1));
                                final long s = Long.parseLong(m2.group(2));
                                String g3 = m2.group(3);
                                while (g3.length() < 3) {
                                    g3 = g3 + "0";
                                }
                                final long millis = Long.parseLong(g3);
                                return m * 60000L + s * 1000L + millis;
                            }
                            return 0L;
                        // linux only; macOS version is "lstart"
                        case "start_time":
                        // see "start_time"
                        case "lstart":
                        // linux only; macOS equivalent is "uid"
                        case "euid":
                        // see "euid"
                        case "uid":
                        default:
                            return value;
                    }
                }
            });
        }

        public RelDataType getRowType(RelDataTypeFactory typeFactory) {
            return typeFactory.builder().add("pid", SqlTypeName.INTEGER).add("ppid", SqlTypeName.INTEGER).add("pgrp", SqlTypeName.INTEGER).add("tpgid", SqlTypeName.INTEGER).add("stat", SqlTypeName.VARCHAR).add("user", SqlTypeName.VARCHAR).add("pcpu", SqlTypeName.DECIMAL, 3, 1).add("pmem", SqlTypeName.DECIMAL, 3, 1).add("vsz", SqlTypeName.INTEGER).add("rss", SqlTypeName.INTEGER).add("tty", SqlTypeName.VARCHAR).add("start_time", SqlTypeName.VARCHAR).add("time", TimeUnit.HOUR, -1, TimeUnit.SECOND, 0).add("euid", SqlTypeName.VARCHAR).add("ruid", SqlTypeName.VARCHAR).add("sess", SqlTypeName.VARCHAR).add("command", SqlTypeName.VARCHAR).build();
        }

        public Statistic getStatistic() {
            return Statistics.of(1000d, ImmutableList.of(ImmutableBitSet.of(1)));
        }

        public Schema.TableType getJdbcTableType() {
            return Schema.TableType.TABLE;
        }

        public boolean isRolledUp(String column) {
            return false;
        }

        public boolean rolledUpColumnValidInsideAgg(String column, SqlCall call, SqlNode parent, CalciteConnectionConfig config) {
            return true;
        }
    };
}

Example 33 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project druid by druid-io.

the class SystemSchemaTest method testServersTable.

@Test
public void testServersTable() {
    SystemSchema.ServersTable serversTable = EasyMock.createMockBuilder(SystemSchema.ServersTable.class).withConstructor(druidNodeDiscoveryProvider, serverInventoryView, authMapper, overlordClient, coordinatorClient).createMock();
    EasyMock.replay(serversTable);
    final DruidNodeDiscovery coordinatorNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
    final DruidNodeDiscovery overlordNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
    final DruidNodeDiscovery brokerNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
    final DruidNodeDiscovery routerNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
    final DruidNodeDiscovery historicalNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
    final DruidNodeDiscovery mmNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
    final DruidNodeDiscovery peonNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
    final DruidNodeDiscovery indexerNodeDiscovery = EasyMock.createMock(DruidNodeDiscovery.class);
    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.COORDINATOR)).andReturn(coordinatorNodeDiscovery).once();
    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.OVERLORD)).andReturn(overlordNodeDiscovery).once();
    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.BROKER)).andReturn(brokerNodeDiscovery).once();
    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.ROUTER)).andReturn(routerNodeDiscovery).once();
    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.HISTORICAL)).andReturn(historicalNodeDiscovery).once();
    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.MIDDLE_MANAGER)).andReturn(mmNodeDiscovery).once();
    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.INDEXER)).andReturn(indexerNodeDiscovery).once();
    EasyMock.expect(druidNodeDiscoveryProvider.getForNodeRole(NodeRole.PEON)).andReturn(peonNodeDiscovery).once();
    EasyMock.expect(coordinatorNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(coordinator, coordinator2)).once();
    EasyMock.expect(overlordNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(overlord, overlord2)).once();
    EasyMock.expect(brokerNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(broker1, broker2, brokerWithBroadcastSegments)).once();
    EasyMock.expect(routerNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(router)).once();
    EasyMock.expect(historicalNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(historical1, historical2, lameHistorical)).once();
    EasyMock.expect(mmNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(middleManager)).once();
    EasyMock.expect(peonNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(peon1, peon2)).once();
    EasyMock.expect(indexerNodeDiscovery.getAllNodes()).andReturn(ImmutableList.of(indexer)).once();
    EasyMock.expect(coordinatorClient.findCurrentLeader()).andReturn(coordinator.getDruidNode().getHostAndPortToUse()).once();
    EasyMock.expect(overlordClient.findCurrentLeader()).andReturn(overlord.getDruidNode().getHostAndPortToUse()).once();
    final List<DruidServer> servers = new ArrayList<>();
    servers.add(mockDataServer(historical1.getDruidNode().getHostAndPortToUse(), 200L, 1000L, "tier"));
    servers.add(mockDataServer(historical2.getDruidNode().getHostAndPortToUse(), 400L, 1000L, "tier"));
    servers.add(mockDataServer(peon1.getDruidNode().getHostAndPortToUse(), 0L, 1000L, "tier"));
    servers.add(mockDataServer(peon2.getDruidNode().getHostAndPortToUse(), 0L, 1000L, "tier"));
    servers.add(mockDataServer(broker1.getDruidNode().getHostAndPortToUse(), 0L, 1000L, "tier"));
    servers.add(mockDataServer(broker2.getDruidNode().getHostAndPortToUse(), 0L, 1000L, "tier"));
    servers.add(mockDataServer(indexer.getDruidNode().getHostAndPortToUse(), 0L, 1000L, "tier"));
    servers.add(mockDataServer(brokerWithBroadcastSegments.getDruidNode().getHostAndPortToUse(), 0L, 1000L, "tier"));
    EasyMock.expect(serverInventoryView.getInventoryValue(lameHistorical.getDruidNode().getHostAndPortToUse())).andReturn(null).once();
    EasyMock.replay(druidNodeDiscoveryProvider, serverInventoryView, coordinatorClient, overlordClient);
    EasyMock.replay(servers.toArray(new Object[0]));
    EasyMock.replay(coordinatorNodeDiscovery, overlordNodeDiscovery, brokerNodeDiscovery, routerNodeDiscovery, historicalNodeDiscovery, mmNodeDiscovery, peonNodeDiscovery, indexerNodeDiscovery);
    DataContext dataContext = createDataContext(Users.SUPER);
    final List<Object[]> rows = serversTable.scan(dataContext).toList();
    rows.sort((Object[] row1, Object[] row2) -> ((Comparable) row1[0]).compareTo(row2[0]));
    final List<Object[]> expectedRows = new ArrayList<>();
    final Long nonLeader = NullHandling.defaultLongValue();
    expectedRows.add(createExpectedRow("brokerHost:8082", "brokerHost", 8082, -1, NodeRole.BROKER, null, 0L, 0L, nonLeader));
    expectedRows.add(createExpectedRow("brokerHostWithBroadcastSegments:8282", "brokerHostWithBroadcastSegments", 8082, 8282, NodeRole.BROKER, "tier", 0L, 1000L, nonLeader));
    expectedRows.add(createExpectedRow("histHost:8083", "histHost", 8083, -1, NodeRole.HISTORICAL, "tier", 400L, 1000L, nonLeader));
    expectedRows.add(createExpectedRow("indexerHost:8091", "indexerHost", 8091, -1, NodeRole.INDEXER, "tier", 0L, 1000L, nonLeader));
    expectedRows.add(createExpectedRow("lameHost:8083", "lameHost", 8083, -1, NodeRole.HISTORICAL, "tier", 0L, 1000L, nonLeader));
    expectedRows.add(createExpectedRow("localhost:8080", "localhost", 8080, -1, NodeRole.PEON, "tier", 0L, 1000L, nonLeader));
    expectedRows.add(createExpectedRow("localhost:8081", "localhost", 8081, -1, NodeRole.COORDINATOR, null, 0L, 0L, 1L));
    expectedRows.add(createExpectedRow("localhost:8082", "localhost", 8082, -1, NodeRole.BROKER, null, 0L, 0L, nonLeader));
    expectedRows.add(createExpectedRow("localhost:8083", "localhost", 8083, -1, NodeRole.HISTORICAL, "tier", 200L, 1000L, nonLeader));
    expectedRows.add(createExpectedRow("localhost:8090", "localhost", 8090, -1, NodeRole.OVERLORD, null, 0L, 0L, 1L));
    expectedRows.add(createExpectedRow("localhost:8181", "localhost", 8181, -1, NodeRole.COORDINATOR, null, 0L, 0L, 0L));
    expectedRows.add(createExpectedRow("localhost:8190", "localhost", 8190, -1, NodeRole.OVERLORD, null, 0L, 0L, 0L));
    expectedRows.add(createExpectedRow("localhost:8888", "localhost", 8888, -1, NodeRole.ROUTER, null, 0L, 0L, nonLeader));
    expectedRows.add(createExpectedRow("mmHost:8091", "mmHost", 8091, -1, NodeRole.MIDDLE_MANAGER, null, 0L, 0L, nonLeader));
    expectedRows.add(createExpectedRow("peonHost:8080", "peonHost", 8080, -1, NodeRole.PEON, "tier", 0L, 1000L, nonLeader));
    Assert.assertEquals(expectedRows.size(), rows.size());
    for (int i = 0; i < rows.size(); i++) {
        Assert.assertArrayEquals(expectedRows.get(i), rows.get(i));
    }
    // Verify value types.
    verifyTypes(rows, SystemSchema.SERVERS_SIGNATURE);
}

Example 34 with DataContext

use of org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.DataContext in project drill by apache.

the class EnumerableRecordReader method setup.

@SuppressWarnings("unchecked")
private void setup(OperatorContext context) {
    SchemaPlus rootSchema = context.getFragmentContext().getFullRootSchema();
    DataContext root = new DrillDataContext(schemaPath != null ? SchemaUtilites.searchSchemaTree(rootSchema, SchemaUtilites.getSchemaPathAsList(schemaPath)) : rootSchema, new JavaTypeFactoryImpl(), Collections.emptyMap());
    try {
        Class<?> implementationClass = ClassBuilder.getCompiledClass(code, CLASS_NAME, context.getFragmentContext().getConfig(), context.getFragmentContext().getOptions());
        Iterable<?> iterable = (Iterable<Map<String, Object>>) implementationClass.getMethod(BuiltInMethod.BINDABLE_BIND.method.getName(), DataContext.class).invoke(implementationClass.newInstance(), root);
        if (fieldsMap.keySet().size() == 1) {
            // for the case of projecting single column, its value is returned
            records = StreamSupport.stream(iterable.spliterator(), false).map(this::wrap).iterator();
        } else {
            // for the case when all columns were projected, array is returned
            records = StreamSupport.stream(iterable.spliterator(), false).map(row -> wrap((Object[]) row)).iterator();
        }
    } catch (CompileException | IOException | ClassTransformationException | ReflectiveOperationException e) {
        logger.error("Exception happened when executing generated code", e);
        Throwable rootCause = Throwables.getRootCause(e);
        throw new DrillRuntimeException(rootCause.getMessage(), rootCause);
    }
}