Java Code Examples for org.apache.calcite.util.Pair#getKey()

The following examples show how to use org.apache.calcite.util.Pair#getKey() . You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may check out the related API usage on the sidebar.
Example 1
Source File: MoreRelOptUtil.java    From dremio-oss with Apache License 2.0 6 votes vote down vote up
public static boolean isSimpleColumnSelection(Project project) {
  HashSet<Integer> inputRefReferenced = new HashSet<>();
  for (Pair<RexNode, String> proj : project.getNamedProjects()) {
    if (proj.getKey().getKind() != SqlKind.INPUT_REF) {
      return false;
    }
    RexInputRef inputRef = (RexInputRef) proj.getKey();
    // If the input reference is again referenced, then it is not a simple column selection (since it is not a permutation).
    if (inputRefReferenced.contains(inputRef.getIndex())) {
      return false;
    }
    final String nameOfProjectField = proj.getValue();
    final String nameOfInput = project.getInput().getRowType().getFieldNames().get(inputRef.getIndex());
    // Renaming a column is not a simple column selection
    if (nameOfProjectField == null || !nameOfProjectField.equals(nameOfInput)) {
      return false;
    }
    inputRefReferenced.add(inputRef.getIndex());
  }
  return true;
}
 
Example 2
Source File: VisitorDataContext.java    From calcite with Apache License 2.0 6 votes vote down vote up
public static DataContext of(RelDataType rowType,
    List<Pair<RexInputRef, RexNode>> usageList) {
  final int size = rowType.getFieldList().size();
  final Object[] values = new Object[size];
  for (Pair<RexInputRef, RexNode> elem : usageList) {
    Pair<Integer, ?> value = getValue(elem.getKey(), elem.getValue());
    if (value == null) {
      LOGGER.warn("{} is not handled for {} for checking implication",
          elem.getKey(), elem.getValue());
      return null;
    }
    int index = value.getKey();
    values[index] = value.getValue();
  }
  return new VisitorDataContext(values);
}
 
Example 3
Source File: SolrFilter.java    From lucene-solr with Apache License 2.0 6 votes vote down vote up
private String translateComparison(RexNode node) {
  Pair<String, RexLiteral> binaryTranslated = null;
  if (((RexCall) node).getOperands().size() == 2) {
    binaryTranslated = translateBinary((RexCall) node);
  }

  switch (node.getKind()) {
    case EQUALS:
      String terms = binaryTranslated.getValue().toString().trim();
      String clause = "eq(" + binaryTranslated.getKey() + "," + terms + ")";
      return clause;
    case NOT_EQUALS:
      return "not(eq(" + binaryTranslated.getKey() + "," + binaryTranslated.getValue() + "))";
    case LESS_THAN:
      return "lt(" + binaryTranslated.getKey() + "," + binaryTranslated.getValue() + ")";
    case LESS_THAN_OR_EQUAL:
      return "lteq(" + binaryTranslated.getKey() + "," + binaryTranslated.getValue() + ")";
    case GREATER_THAN:
      return "gt(" + binaryTranslated.getKey() + "," + binaryTranslated.getValue() + ")";
    case GREATER_THAN_OR_EQUAL:
      return "gteq(" + binaryTranslated.getKey() + "," + binaryTranslated.getValue() + ")";
    default:
      throw new AssertionError("cannot translate " + node);
  }
}
 
Example 4
Source File: SolrTable.java    From lucene-solr with Apache License 2.0 6 votes vote down vote up
private Metric getMetric(Pair<String, String> metricPair) {
  switch (metricPair.getKey()) {
    case "COUNT":
      return new CountMetric(metricPair.getValue());
    case "SUM":
    case "$SUM0":
      return new SumMetric(metricPair.getValue());
    case "MIN":
      return new MinMetric(metricPair.getValue());
    case "MAX":
      return new MaxMetric(metricPair.getValue());
    case "AVG":
      return new MeanMetric(metricPair.getValue());
    default:
      throw new IllegalArgumentException(metricPair.getKey());
  }
}
 
Example 5
Source File: NormalHandler.java    From dremio-oss with Apache License 2.0 6 votes vote down vote up
@Override
public PhysicalPlan getPlan(SqlHandlerConfig config, String sql, SqlNode sqlNode) throws Exception {
  try{
    final ConvertedRelNode convertedRelNode = PrelTransformer.validateAndConvert(config, sqlNode);
    final RelDataType validatedRowType = convertedRelNode.getValidatedRowType();
    final RelNode queryRelNode = convertedRelNode.getConvertedNode();

    final Rel drel = PrelTransformer.convertToDrel(config, queryRelNode, validatedRowType);

    final Pair<Prel, String> convertToPrel = PrelTransformer.convertToPrel(config, drel);
    final Prel prel = convertToPrel.getKey();
    textPlan = convertToPrel.getValue();
    final PhysicalOperator pop = PrelTransformer.convertToPop(config, prel);
    final PhysicalPlan plan = PrelTransformer.convertToPlan(config, pop);
    logger.debug("Final Physical Plan {}", textPlan);
    PrelTransformer.log(config, "Dremio Plan", plan, logger);


    return plan;
  }catch(Exception ex){
    throw SqlExceptionHelper.coerceException(logger, sql, ex, true);
  }
}
 
Example 6
Source File: ClassGenerator.java    From Bats with Apache License 2.0 6 votes vote down vote up
/**
 * declare a constant field for the class.
 * argument {@code function} holds the constant value which
 * returns a value holder must be set to the class field when the class instance created.
 * the class field innerClassField will be created if innerClassGenerator exists.
 *
 * @param prefix the prefix name of class field
 * @param t the type of class field
 * @param init init expression
 * @param function the function holds the constant value
 * @return the depth of nested class, class field
 */
public Pair<Integer, JVar> declareClassConstField(String prefix, JType t, JExpression init,
                                                  Function<DrillBuf, ? extends ValueHolder> function) {
  JVar var;
  int depth = 1;
  if (innerClassGenerator != null) {
    Pair<Integer, JVar> nested = innerClassGenerator.declareClassConstField(prefix, t, init, function);
    depth = nested.getKey() + 1;
    var = nested.getValue();
  } else {
    var = clazz.field(JMod.NONE, t, prefix + index++, init);
  }
  Pair<Integer, JVar> depthVar = Pair.of(depth, var);
  constantVars.put(depthVar, function);
  return depthVar;
}
 
Example 7
Source File: VisitorDataContext.java    From Bats with Apache License 2.0 6 votes vote down vote up
public static DataContext of(RelDataType rowType,
    List<Pair<RexInputRef, RexNode>> usageList) {
  final int size = rowType.getFieldList().size();
  final Object[] values = new Object[size];
  for (Pair<RexInputRef, RexNode> elem : usageList) {
    Pair<Integer, ?> value = getValue(elem.getKey(), elem.getValue());
    if (value == null) {
      LOGGER.warn("{} is not handled for {} for checking implication",
          elem.getKey(), elem.getValue());
      return null;
    }
    int index = value.getKey();
    values[index] = value.getValue();
  }
  return new VisitorDataContext(values);
}
 
Example 8
Source File: PrelTransformer.java    From dremio-oss with Apache License 2.0 5 votes vote down vote up
public static ConvertedRelNode validateAndConvert(SqlHandlerConfig config, SqlNode sqlNode, RelTransformer relTransformer) throws ForemanSetupException, RelConversionException, ValidationException {
  final Pair<SqlNode, RelDataType> validatedTypedSqlNode = validateNode(config, sqlNode);
  if (config.getObserver() != null) {
    config.getObserver().beginState(AttemptObserver.toEvent(UserBitShared.AttemptEvent.State.PLANNING));
  }

  final SqlNode validated = validatedTypedSqlNode.getKey();
  final RelNode rel = convertToRel(config, validated, relTransformer);
  final RelNode preprocessedRel = preprocessNode(config, rel);
  return new ConvertedRelNode(preprocessedRel, validatedTypedSqlNode.getValue());
}
 
Example 9
Source File: VisitorDataContext.java    From Bats with Apache License 2.0 5 votes vote down vote up
public static DataContext of(RelDataType rowType, RexNode rex) {
  final int size = rowType.getFieldList().size();
  final Object[] values = new Object[size];
  final List<RexNode> operands = ((RexCall) rex).getOperands();
  final RexNode firstOperand = operands.get(0);
  final RexNode secondOperand = operands.get(1);
  final Pair<Integer, ?> value = getValue(firstOperand, secondOperand);
  if (value != null) {
    int index = value.getKey();
    values[index] = value.getValue();
    return new VisitorDataContext(values);
  } else {
    return null;
  }
}
 
Example 10
Source File: SolrTable.java    From lucene-solr with Apache License 2.0 5 votes vote down vote up
private static FieldComparator[] getComps(List<Pair<String, String>> orders) {
  FieldComparator[] comps = new FieldComparator[orders.size()];
  for(int i=0; i<orders.size(); i++) {
    Pair<String,String> sortItem = orders.get(i);
    String ordering = sortItem.getValue();
    ComparatorOrder comparatorOrder = ascDescComp(ordering);
    String sortKey = sortItem.getKey();
    comps[i] = new FieldComparator(sortKey, comparatorOrder);
  }

  return comps;
}
 
Example 11
Source File: SolrFilter.java    From lucene-solr with Apache License 2.0 5 votes vote down vote up
private String translateComparison(RexNode node) {
  Pair<String, RexLiteral> binaryTranslated = null;
  if (((RexCall) node).getOperands().size() == 2) {
    binaryTranslated = translateBinary((RexCall) node);
  }

  switch (node.getKind()) {
    case NOT:
      return "-" + translateComparison(((RexCall) node).getOperands().get(0));
    case EQUALS:
      String terms = binaryTranslated.getValue().toString().trim();
      terms = terms.replace("'","");
      if (!terms.startsWith("(") && !terms.startsWith("[") && !terms.startsWith("{")) {
        terms = "\"" + terms + "\"";
      }

      String clause = binaryTranslated.getKey() + ":" + terms;
      this.negativeQuery = false;
      return clause;
    case NOT_EQUALS:
      return "-(" + binaryTranslated.getKey() + ":" + binaryTranslated.getValue() + ")";
    case LESS_THAN:
      this.negativeQuery = false;
      return "(" + binaryTranslated.getKey() + ": [ * TO " + binaryTranslated.getValue() + " })";
    case LESS_THAN_OR_EQUAL:
      this.negativeQuery = false;
      return "(" + binaryTranslated.getKey() + ": [ * TO " + binaryTranslated.getValue() + " ])";
    case GREATER_THAN:
      this.negativeQuery = false;
      return "(" + binaryTranslated.getKey() + ": { " + binaryTranslated.getValue() + " TO * ])";
    case GREATER_THAN_OR_EQUAL:
      this.negativeQuery = false;
      return "(" + binaryTranslated.getKey() + ": [ " + binaryTranslated.getValue() + " TO * ])";
    default:
      throw new AssertionError("cannot translate " + node);
  }
}
 
Example 12
Source File: DefaultSqlHandler.java    From Bats with Apache License 2.0 5 votes vote down vote up
protected ConvertedRelNode validateAndConvert(SqlNode sqlNode) throws ForemanSetupException, RelConversionException, ValidationException {
  final SqlNode rewrittenSqlNode = rewrite(sqlNode);
  final Pair<SqlNode, RelDataType> validatedTypedSqlNode = validateNode(rewrittenSqlNode);
  final SqlNode validated = validatedTypedSqlNode.getKey();

  RelNode rel = convertToRel(validated);
  rel = preprocessNode(rel);

  return new ConvertedRelNode(rel, validatedTypedSqlNode.getValue());
}
 
Example 13
Source File: SolrAggregate.java    From lucene-solr with Apache License 2.0 5 votes vote down vote up
public void implement(Implementor implementor) {
  implementor.visitChild(0, getInput());

  final List<String> inNames = SolrRules.solrFieldNames(getInput().getRowType());


  for(Pair<AggregateCall, String> namedAggCall : getNamedAggCalls()) {

    AggregateCall aggCall = namedAggCall.getKey();

    Pair<String, String> metric = toSolrMetric(implementor, aggCall, inNames);
    implementor.addReverseAggMapping(namedAggCall.getValue(), metric.getKey().toLowerCase(Locale.ROOT)+"("+metric.getValue()+")");
    implementor.addMetricPair(namedAggCall.getValue(), metric.getKey(), metric.getValue());
    /*
    if(aggCall.getName() == null) {
      System.out.println("AGG:"+namedAggCall.getValue()+":"+ aggCall.getAggregation().getName() + "(" + inNames.get(aggCall.getArgList().get(0)) + ")");
      implementor.addFieldMapping(namedAggCall.getValue(),
        aggCall.getAggregation().getName() + "(" + inNames.get(aggCall.getArgList().get(0)) + ")");
    }
    */
  }

  for(int group : getGroupSet()) {
    String inName = inNames.get(group);
    implementor.addBucket(inName);
  }
}
 
Example 14
Source File: VisitorDataContext.java    From calcite with Apache License 2.0 5 votes vote down vote up
public static DataContext of(RelDataType rowType, RexNode rex) {
  final int size = rowType.getFieldList().size();
  final Object[] values = new Object[size];
  final List<RexNode> operands = ((RexCall) rex).getOperands();
  final RexNode firstOperand = operands.get(0);
  final RexNode secondOperand = operands.get(1);
  final Pair<Integer, ?> value = getValue(firstOperand, secondOperand);
  if (value != null) {
    int index = value.getKey();
    values[index] = value.getValue();
    return new VisitorDataContext(values);
  } else {
    return null;
  }
}
 
Example 15
Source File: JournalledUpdateRule.java    From calcite-sql-rewriter with Apache License 2.0 4 votes vote down vote up
@Override
public RelNode doApply(LogicalTableModify tableModify, JournalledJdbcTable journalTable,
		JdbcRelBuilderFactory relBuilderFactory) {

	if (!(tableModify.getInput() instanceof LogicalProject)) {
		throw new IllegalStateException("Unknown Calcite UPDATE structure");
	}

	String versionField = journalTable.getVersionField();

	// Merge the Update's update column expression into the target INSERT
	LogicalProject project = (LogicalProject) tableModify.getInput();
	List<RexNode> desiredFields = new ArrayList<>();
	List<String> desiredNames = new ArrayList<>();

	for (Pair<RexNode, String> field : project.getNamedProjects()) {
		if (field.getKey() instanceof RexInputRef) {
			int index = tableModify.getUpdateColumnList().indexOf(field.getValue());
			if (index != -1) {
				desiredFields.add(tableModify.getSourceExpressionList().get(index));
			}
			else {
				desiredFields.add(field.getKey());
			}
			desiredNames.add(field.getValue());
		}
	}

	JdbcRelBuilder relBuilder = relBuilderFactory.create(
			tableModify.getCluster(),
			tableModify.getTable().getRelOptSchema()
	);

	relBuilder.push(project.getInput());

	JournalVersionType versionType = journalTable.getVersionType();
	if (!versionType.isValidSqlType(relBuilder.field(versionField).getType().getSqlTypeName())) {
		throw new IllegalStateException("Incorrect journalVersionType! Column 'version_number' is of type: "
				+ relBuilder.field(versionField).getType().getSqlTypeName()
				+ " but the journalVersionType is " + versionType);
	}
	if (versionType.updateRequiresExplicitVersion()) {
		RexNode newVersion = versionType.incrementVersion(relBuilder, relBuilder.field(versionField));
		desiredFields.add(newVersion);
		desiredNames.add(versionField);
	}

	relBuilder.project(desiredFields, desiredNames);

	// Convert the UPDATE into INSERT TableModify operations
	relBuilder.insertCopying(
			tableModify,
			journalTable.getJournalTable()
	);

	return relBuilder.build();
}
 
Example 16
Source File: DataAdditionCmdHandler.java    From dremio-oss with Apache License 2.0 4 votes vote down vote up
public PhysicalPlan getPlan(DatasetCatalog datasetCatalog, NamespaceKey path, SqlHandlerConfig config, String sql, SqlNode sqlNode, DataAdditionCmdCall sqlCmd) throws Exception {
  try {
    final ConvertedRelNode convertedRelNode = PrelTransformer.validateAndConvert(config, sqlCmd.getQuery());
    final RelDataType validatedRowType = convertedRelNode.getValidatedRowType();

    long maxColumnCount = config.getContext().getOptions().getOption(CatalogOptions.METADATA_LEAF_COLUMN_MAX);
    if (validatedRowType.getFieldCount() > maxColumnCount) {
      throw new ColumnCountTooLargeException((int) maxColumnCount);
    }

    final RelNode queryRelNode = convertedRelNode.getConvertedNode();
    final RelNode newTblRelNode = SqlHandlerUtil.resolveNewTableRel(false, sqlCmd.getFieldNames(),
        validatedRowType, queryRelNode, !isCreate());

    final long ringCount = config.getContext().getOptions().getOption(PlannerSettings.RING_COUNT);

    // For Insert command we make sure query schema match exactly with table schema,
    // which includes partition columns. So, not checking here
    final RelNode newTblRelNodeWithPCol = SqlHandlerUtil.qualifyPartitionCol(newTblRelNode,
      isCreate() ? sqlCmd.getPartitionColumns(datasetCatalog, path) : Lists.newArrayList());

    PrelTransformer.log("Calcite", newTblRelNodeWithPCol, logger, null);

    ByteString extendedByteString = null;
    if (!isCreate()) {
      DremioTable table = datasetCatalog.getTable(path);
      extendedByteString = table.getDatasetConfig().getReadDefinition().getExtendedProperty();
    }

    final WriterOptions options = new WriterOptions(
      (int) ringCount,
      sqlCmd.getPartitionColumns(datasetCatalog, path),
      sqlCmd.getSortColumns(),
      sqlCmd.getDistributionColumns(),
      sqlCmd.getPartitionDistributionStrategy(),
      sqlCmd.isSingleWriter(),
      Long.MAX_VALUE,
      getIcebergWriterOperation(),
      extendedByteString
    );

    // Convert the query to Dremio Logical plan and insert a writer operator on top.
    Rel drel = this.convertToDrel(
      config,
      newTblRelNodeWithPCol,
      datasetCatalog,
      path,
      options,
      newTblRelNode.getRowType(),
      storageOptionsMap, sqlCmd.getFieldNames());

    final Pair<Prel, String> convertToPrel = PrelTransformer.convertToPrel(config, drel);
    final Prel prel = convertToPrel.getKey();
    textPlan = convertToPrel.getValue();
    PhysicalOperator pop = PrelTransformer.convertToPop(config, prel);

    PhysicalPlan plan = PrelTransformer.convertToPlan(config, pop,
      isIcebergTable() ? () -> refreshDataset(datasetCatalog, path, isCreate()) : null);

    PrelTransformer.log(config, "Dremio Plan", plan, logger);

    return plan;

  } catch(Exception ex){
    throw SqlExceptionHelper.coerceException(logger, sql, ex, true);
  }
}
 
Example 17
Source File: JournalledInsertRule.java    From calcite-sql-rewriter with Apache License 2.0 4 votes vote down vote up
@Override
public RelNode doApply(LogicalTableModify tableModify, JournalledJdbcTable journalTable,
		JdbcRelBuilderFactory relBuilderFactory) {

	JdbcRelBuilder relBuilder = relBuilderFactory.create(
			tableModify.getCluster(),
			tableModify.getTable().getRelOptSchema()
	);

	RelNode input = tableModify.getInput();
	if (input instanceof LogicalValues) {

		// TODO: do we need to do anything here?
		relBuilder.push(input);

	}
	else if (input instanceof LogicalProject) {

		LogicalProject project = (LogicalProject) input;
		List<RexNode> desiredFields = new ArrayList<>();
		List<String> desiredNames = new ArrayList<>();
		for (Pair<RexNode, String> field : project.getNamedProjects()) {
			if (field.getKey() instanceof RexInputRef) {
				desiredFields.add(field.getKey());
				desiredNames.add(field.getValue());
			}
		}

		relBuilder.push(project.getInput());
		relBuilder.project(desiredFields, desiredNames);

	}
	else {
		throw new IllegalStateException("Unknown Calcite INSERT structure");
	}

	relBuilder.insertCopying(
			tableModify,
			journalTable.getJournalTable()
	);

	return relBuilder.build();

}
 
Example 18
Source File: Limit0LogicalToPhysicalTest.java    From dremio-oss with Apache License 2.0 4 votes vote down vote up
@Ignore
public void ExchangesKeepTest() throws Exception {

  final String yelpTable = TEMP_SCHEMA + ".\"yelp\"";
  final String sql = "SELECT nested_0.review_id AS review_id, nested_0.user_id AS user_id, nested_0.votes AS votes," +
    " nested_0.stars AS stars, join_business.business_id AS business_id0, join_business.neighborhoods AS neighborhoods, join_business.city AS city, join_business.latitude AS latitude, join_business.review_count AS review_count, join_business.full_address AS full_address, join_business.stars AS stars0, join_business.categories AS categories, join_business.state AS state, join_business.longitude AS longitude\n" +
    "FROM (\n" +
    "  SELECT review_id, user_id, votes, stars, business_id\n" +
    "  FROM cp.\"yelp_review.json\" where 1 = 0\n" +
    ") nested_0\n" +
    " FULL JOIN " + yelpTable + " AS join_business ON nested_0.business_id = join_business.business_id";

  final SabotContext context = getSabotContext();
  context.getOptionManager().setOption(
    OptionValue.createLong(OptionValue.OptionType.SYSTEM, "planner.slice_target", 1)
  );
  context.getOptionManager().setOption(
    OptionValue.createLong(OptionValue.OptionType.SYSTEM, "planner.width.max_per_node", 10)
  );
  context.getOptionManager().setOption(
    OptionValue.createBoolean(OptionValue.OptionType.SYSTEM, "planner.enable_mux_exchange", true)
  );

  final QueryContext queryContext = new QueryContext(session(), context, UserBitShared.QueryId.getDefaultInstance());
  final AttemptObserver observer = new PassthroughQueryObserver(ExecTest.mockUserClientConnection(null));
  final SqlConverter converter = new SqlConverter(
      queryContext.getPlannerSettings(),
      queryContext.getOperatorTable(),
      queryContext,
      queryContext.getMaterializationProvider(),
      queryContext.getFunctionRegistry(),
      queryContext.getSession(),
      observer,
      queryContext.getCatalog(),
      queryContext.getSubstitutionProviderFactory(),
      queryContext.getConfig(),
      queryContext.getScanResult());
  final SqlNode node = converter.parse(sql);
  final SqlHandlerConfig config = new SqlHandlerConfig(queryContext, converter, observer, null);

  final ConvertedRelNode convertedRelNode = PrelTransformer.validateAndConvert(config, node);
  final RelDataType validatedRowType = convertedRelNode.getValidatedRowType();
  final RelNode queryRelNode = convertedRelNode.getConvertedNode();

  final Rel drel = PrelTransformer.convertToDrel(config, queryRelNode, validatedRowType);

  final Pair<Prel, String> convertToPrel = PrelTransformer.convertToPrel(config, drel);
  final Prel prel = convertToPrel.getKey();
  final String prePhysicaltextPlan = convertToPrel.getValue();

  assertThat(prePhysicaltextPlan, CoreMatchers.containsString("HashToRandomExchange"));
  assertThat(prePhysicaltextPlan, CoreMatchers.containsString("UnorderedMuxExchange"));
  assertThat(prePhysicaltextPlan, CoreMatchers.containsString("Empty"));
  assertThat(prePhysicaltextPlan, CoreMatchers.containsString("EasyScan"));

  final PhysicalOperator pop = PrelTransformer.convertToPop(config, prel);
  final PhysicalPlan plan = PrelTransformer.convertToPlan(config, pop);
  final String postPhysicaltextPlan = plan.unparse(config.getContext().getLpPersistence().getMapper().writer());

  assertThat(postPhysicaltextPlan, CoreMatchers.containsString("EmptyValues"));
  assertThat(postPhysicaltextPlan, CoreMatchers.containsString("EasyGroupScan"));
  assertThat(postPhysicaltextPlan, CoreMatchers.containsString("unordered-mux-exchange"));
  assertThat(postPhysicaltextPlan, CoreMatchers.containsString("hash-to-random-exchange"));


  PhysicalPlanReader pPlanReader = new PhysicalPlanReader(
    DEFAULT_SABOT_CONFIG, CLASSPATH_SCAN_RESULT, new LogicalPlanPersistence(DEFAULT_SABOT_CONFIG, CLASSPATH_SCAN_RESULT),
    CoordinationProtos.NodeEndpoint.getDefaultInstance(),
    DirectProvider.wrap(Mockito.mock(CatalogService.class)), context);

  ExecutionPlan exec = ExecutionPlanCreator
    .getExecutionPlan(queryContext, pPlanReader, AbstractMaestroObserver.NOOP, plan,
    QueueType.SMALL);
  List<PlanFragmentFull> fragments  = exec.getFragments();

  int scanFrags = 0;
  for (PlanFragmentFull fragment : fragments) {
    if (new String(fragment.getMajor().getFragmentJson().toByteArray()).contains("easy-sub-scan")) {
      scanFrags++;
    }
  }
  assertEquals(2, scanFrags);

}
 
Example 19
Source File: CompactRefreshHandler.java    From dremio-oss with Apache License 2.0 4 votes vote down vote up
@Override
public PhysicalPlan getPlan(SqlHandlerConfig config, String sql, SqlNode sqlNode) throws Exception {
  try {
    final SqlCompactMaterialization compact = SqlNodeUtil.unwrap(sqlNode, SqlCompactMaterialization.class);

    if(!SystemUser.SYSTEM_USERNAME.equals(config.getContext().getQueryUserName())) {
      throw SqlExceptionHelper.parseError("$COMPACT REFRESH not supported.", sql, compact.getParserPosition())
        .build(logger);
    }

    final AttemptObserver observer = config.getObserver();
    ReflectionService service = config.getContext().getAccelerationManager().unwrap(ReflectionService.class);

    // Let's validate the plan.
    Stopwatch watch = Stopwatch.createStarted();

    final List<String> materializationPath = normalizeComponents(compact.getMaterializationPath());
    if (materializationPath == null) {
      throw SqlExceptionHelper.parseError("Unknown materialization", sql, compact.getParserPosition())
        .build(logger);
    }

    final ReflectionId reflectionId = new ReflectionId(materializationPath.get(0));
    Optional<ReflectionGoal> goalOpt = service.getGoal(reflectionId);
    if(!goalOpt.isPresent()) {
      throw SqlExceptionHelper.parseError("Unknown reflection id.", sql, compact.getParserPosition()).build(logger);
    }
    final ReflectionGoal goal = goalOpt.get();

    Optional<ReflectionEntry> entryOpt = service.getEntry(reflectionId);
    if(!entryOpt.isPresent()) {
      throw SqlExceptionHelper.parseError("Unknown reflection id.", sql, compact.getParserPosition()).build(logger);
    }
    final ReflectionEntry entry = entryOpt.get();
    if(!ReflectionGoalsStore.checkGoalVersion(goal, entry.getGoalVersion())) {
      throw UserException.validationError().message("Reflection has been updated since reflection was scheduled.").build(logger);
    }

    Optional<Materialization> materializationOpt = service.getMaterialization(new MaterializationId(materializationPath.get(1)));
    if (!materializationOpt.isPresent()) {
      throw SqlExceptionHelper.parseError("Unknown materialization id", sql, compact.getParserPosition()).build(logger);
    }
    final Materialization materialization = materializationOpt.get();

    List<Refresh> refreshes = Lists.newArrayList(service.getRefreshes(materialization));
    if (refreshes.size() != 1) {
      throw SqlExceptionHelper.parseError("Invalid materialization", sql, compact.getParserPosition()).build(logger);
    }

    Optional<Materialization> newMaterializationOpt = service.getMaterialization(new MaterializationId(compact.getNewMaterializationId()));
    if (!newMaterializationOpt.isPresent()) {
      throw SqlExceptionHelper.parseError("Unknown new materialization id", sql, compact.getParserPosition()).build(logger);
    }
    final Materialization newMaterialization = newMaterializationOpt.get();

    observer.planValidated(CalciteArrowHelper.wrap(RecordWriter.SCHEMA).toCalciteRecordType(config.getConverter().getCluster().getTypeFactory()), compact, watch.elapsed(TimeUnit.MILLISECONDS));

    watch.reset();

    final List<String> tableSchemaPath = ReflectionUtils.getMaterializationPath(materialization);

    final PlanNormalizer planNormalizer = new PlanNormalizer(config);
    final RelNode initial = getPlan(config, tableSchemaPath, planNormalizer);

    observer.planConvertedToRel(initial, watch.elapsed(TimeUnit.MILLISECONDS));

    final Rel drel = PrelTransformer.convertToDrel(config, initial);
    final Set<String> fields = ImmutableSet.copyOf(drel.getRowType().getFieldNames());
    final long ringCount = config.getContext().getOptions().getOption(PlannerSettings.RING_COUNT);
    final Rel writerDrel = new WriterRel(drel.getCluster(), drel.getCluster().traitSet().plus(Rel.LOGICAL),
      drel, config.getContext().getCatalog().createNewTable(
        new NamespaceKey(ReflectionUtils.getMaterializationPath(newMaterialization)),
      null, getWriterOptions((int) ringCount, goal, fields), ImmutableMap.of()),
      initial.getRowType()
    );

    final RelNode doubleWriter = SqlHandlerUtil.storeQueryResultsIfNeeded(config.getConverter().getParserConfig(), config.getContext(), writerDrel);

    final ScreenRel screen = new ScreenRel(writerDrel.getCluster(), writerDrel.getTraitSet(), doubleWriter);

    final Pair<Prel, String> convertToPrel = PrelTransformer.convertToPrel(config, screen);
    final Prel prel = convertToPrel.getKey();
    this.textPlan = convertToPrel.getValue();
    PhysicalOperator pop = PrelTransformer.convertToPop(config, prel);
    PhysicalPlan plan = PrelTransformer.convertToPlan(config, pop);

    if (logger.isTraceEnabled()) {
      PrelTransformer.log(config, "Dremio Plan", plan, logger);
    }

    return plan;
  } catch (Exception e) {
    throw SqlExceptionHelper.coerceException(logger, sql, e, true);
  }
}