/** * Sets up projection that will transfer all of the columns in batch, and also populate the * partition column based on which partition a record falls into in the partition table * * @param batch * @throws SchemaChangeException */ protected void setupNewSchema(VectorAccessible batch) throws SchemaChangeException { container.clear(); final ErrorCollector collector = new ErrorCollectorImpl(); final List<TransferPair> transfers = Lists.newArrayList(); final ClassGenerator<OrderedPartitionProjector> cg = CodeGenerator.getRoot( OrderedPartitionProjector.TEMPLATE_DEFINITION, context.getFunctionRegistry()); for (VectorWrapper<?> vw : batch) { TransferPair tp = vw.getValueVector().getTransferPair(); transfers.add(tp); container.add(tp.getTo()); } cg.setMappingSet(mainMapping); int count = 0; for (Ordering od : popConfig.getOrderings()) { final LogicalExpression expr = ExpressionTreeMaterializer.materialize( od.getExpr(), batch, collector, context.getFunctionRegistry()); if (collector.hasErrors()) throw new SchemaChangeException( "Failure while materializing expression. " + collector.toErrorString()); cg.setMappingSet(incomingMapping); ClassGenerator.HoldingContainer left = cg.addExpr(expr, false); cg.setMappingSet(partitionMapping); ClassGenerator.HoldingContainer right = cg.addExpr( new ValueVectorReadExpression(new TypedFieldId(expr.getMajorType(), count++)), false); cg.setMappingSet(mainMapping); LogicalExpression fh = FunctionGenerationHelper.getComparator(left, right, context.getFunctionRegistry()); ClassGenerator.HoldingContainer out = cg.addExpr(fh, false); JConditional jc = cg.getEvalBlock()._if(out.getValue().ne(JExpr.lit(0))); if (od.getDirection() == Direction.ASCENDING) { jc._then()._return(out.getValue()); } else { jc._then()._return(out.getValue().minus()); } } cg.getEvalBlock()._return(JExpr.lit(0)); container.add(this.partitionKeyVector); container.buildSchema(batch.getSchema().getSelectionVectorMode()); try { this.projector = context.getImplementationClass(cg); projector.setup( context, batch, this, transfers, partitionVectors, partitions, popConfig.getRef()); } catch (ClassTransformationException | IOException e) { throw new SchemaChangeException("Failure while attempting to load generated class", e); } }
protected void doWork(VectorAccessible batch) { int recordCount = batch.getRecordCount(); AllocationHelper.allocate(partitionKeyVector, recordCount, 50); projector.projectRecords(recordCount, 0); for (VectorWrapper<?> v : container) { ValueVector.Mutator m = v.getValueVector().getMutator(); m.setValueCount(recordCount); } }