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

Example 91 with SchemaPlus

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

the class ZetaSQLQueryPlanner method defaultConfig.

private static FrameworkConfig defaultConfig(JdbcConnection connection, Collection<RuleSet> ruleSets) {
    final CalciteConnectionConfig config = connection.config();
    final SqlParser.ConfigBuilder parserConfig = SqlParser.configBuilder().setQuotedCasing(config.quotedCasing()).setUnquotedCasing(config.unquotedCasing()).setQuoting(config.quoting()).setConformance(config.conformance()).setCaseSensitive(config.caseSensitive());
    final SqlParserImplFactory parserFactory = config.parserFactory(SqlParserImplFactory.class, null);
    if (parserFactory != null) {
        parserConfig.setParserFactory(parserFactory);
    }
    final SchemaPlus schema = connection.getRootSchema();
    final SchemaPlus defaultSchema = connection.getCurrentSchemaPlus();
    final ImmutableList<RelTraitDef> traitDefs = ImmutableList.of(ConventionTraitDef.INSTANCE);
    final CalciteCatalogReader catalogReader = new CalciteCatalogReader(CalciteSchema.from(schema), ImmutableList.of(defaultSchema.getName()), connection.getTypeFactory(), connection.config());
    final SqlOperatorTable opTab0 = connection.config().fun(SqlOperatorTable.class, SqlStdOperatorTable.instance());
    return Frameworks.newConfigBuilder().parserConfig(parserConfig.build()).defaultSchema(defaultSchema).traitDefs(traitDefs).ruleSets(ruleSets.toArray(new RuleSet[0])).costFactory(BeamCostModel.FACTORY).typeSystem(connection.getTypeFactory().getTypeSystem()).operatorTable(SqlOperatorTables.chain(opTab0, catalogReader)).build();
}

Example 92 with SchemaPlus

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

the class BeamZetaSqlCatalogTest method rejectsScalarFunctionImplWithUnsupportedParameterType.

@Test
public void rejectsScalarFunctionImplWithUnsupportedParameterType() throws NoSuchMethodException {
    JdbcConnection jdbcConnection = createJdbcConnection();
    SchemaPlus calciteSchema = jdbcConnection.getCurrentSchemaPlus();
    Method method = TakesArrayTimeFn.class.getMethod("eval", List.class);
    calciteSchema.add("take_array", ScalarFunctionImpl.create(method));
    thrown.expect(UnsupportedOperationException.class);
    thrown.expectMessage("Calcite type TIME not allowed in function take_array");
    BeamZetaSqlCatalog.create(calciteSchema, jdbcConnection.getTypeFactory(), SqlAnalyzer.baseAnalyzerOptions());
}

Example 93 with SchemaPlus

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

the class ThreeTablesSchema method transform.

private RelNode transform(String sql, RuleSet prepareRules) throws Exception {
    final SchemaPlus rootSchema = Frameworks.createRootSchema(true);
    final SchemaPlus defSchema = rootSchema.add("tt", new ThreeTablesSchema());
    final FrameworkConfig config = Frameworks.newConfigBuilder().parserConfig(SqlParser.Config.DEFAULT).defaultSchema(defSchema).traitDefs(ConventionTraitDef.INSTANCE, RelCollationTraitDef.INSTANCE).programs(Programs.of(prepareRules)).build();
    Planner planner = Frameworks.getPlanner(config);
    SqlNode parse = planner.parse(sql);
    SqlNode validate = planner.validate(parse);
    RelRoot planRoot = planner.rel(validate);
    RelNode planBefore = planRoot.rel;
    RelTraitSet desiredTraits = planBefore.getTraitSet().replace(EnumerableConvention.INSTANCE);
    return planner.transform(0, desiredTraits, planBefore);
}

Example 94 with SchemaPlus

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

the class ZetaSqlJavaUdfTypeTest method setUp.

@Before
public void setUp() throws NoSuchMethodException {
    initialize();
    // Register test table.
    JdbcConnection jdbcConnection = JdbcDriver.connect(new ReadOnlyTableProvider("table_provider", ImmutableMap.of("table", table)), PipelineOptionsFactory.create());
    // Register UDFs.
    SchemaPlus schema = jdbcConnection.getCurrentSchemaPlus();
    schema.add("test_boolean", ScalarFunctionImpl.create(BooleanIdentityFn.class.getMethod("eval", Boolean.class)));
    schema.add("test_int64", ScalarFunctionImpl.create(Int64IdentityFn.class.getMethod("eval", Long.class)));
    schema.add("test_string", ScalarFunctionImpl.create(StringIdentityFn.class.getMethod("eval", String.class)));
    schema.add("test_bytes", ScalarFunctionImpl.create(BytesIdentityFn.class.getMethod("eval", byte[].class)));
    schema.add("test_float64", ScalarFunctionImpl.create(DoubleIdentityFn.class.getMethod("eval", Double.class)));
    schema.add("test_date", ScalarFunctionImpl.create(DateIdentityFn.class.getMethod("eval", Date.class)));
    schema.add("test_timestamp", ScalarFunctionImpl.create(TimestampIdentityFn.class.getMethod("eval", Timestamp.class)));
    schema.add("test_array", ScalarFunctionImpl.create(ListIdentityFn.class.getMethod("eval", List.class)));
    schema.add("test_numeric", ScalarFunctionImpl.create(BigDecimalIdentityFn.class.getMethod("eval", BigDecimal.class)));
    this.config = Frameworks.newConfigBuilder(config).defaultSchema(schema).build();
}

Example 95 with SchemaPlus

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

the class SqlCreateFunction method execute.

@Override
public void execute(CalcitePrepare.Context context) {
    final Pair<CalciteSchema, String> pair = SqlDdlNodes.schema(context, true, functionName);
    SchemaPlus schema = pair.left.plus();
    String lastName = pair.right;
    if (!schema.getFunctions(lastName).isEmpty()) {
        throw SqlUtil.newContextException(functionName.getParserPosition(), RESOURCE.internal(String.format("Function %s is already defined.", lastName)));
    }
    JavaUdfLoader udfLoader = new JavaUdfLoader();
    // TODO(BEAM-12355) Support qualified function names.
    List<String> functionPath = ImmutableList.of(lastName);
    if (!(jarPath instanceof SqlCharStringLiteral)) {
        throw SqlUtil.newContextException(jarPath.getParserPosition(), RESOURCE.internal("Jar path is not instanceof SqlCharStringLiteral."));
    }
    String unquotedJarPath = ((SqlCharStringLiteral) jarPath).getNlsString().getValue();
    if (isAggregate) {
        // Try loading the aggregate function just to make sure it exists. LazyAggregateCombineFn will
        // need to fetch it again at runtime.
        udfLoader.loadAggregateFunction(functionPath, unquotedJarPath);
        LazyAggregateCombineFn<?, ?, ?> combineFn = new LazyAggregateCombineFn<>(functionPath, unquotedJarPath);
        schema.add(lastName, combineFn.getUdafImpl());
    } else {
        ScalarFn scalarFn = udfLoader.loadScalarFunction(functionPath, unquotedJarPath);
        Method method = ScalarFnReflector.getApplyMethod(scalarFn);
        Function function = ScalarFunctionImpl.create(method, unquotedJarPath);
        schema.add(lastName, function);
    }
}