protected QueryContext mockQueryContext(SabotContext dbContext) throws Exception { final UserSession userSession = UserSession.Builder.newBuilder().withOptionManager(dbContext.getOptionManager()).build(); final SessionOptionManager sessionOptions = (SessionOptionManager) userSession.getOptions(); final QueryOptionManager queryOptions = new QueryOptionManager(sessionOptions); final ExecutionControls executionControls = new ExecutionControls(queryOptions, NodeEndpoint.getDefaultInstance()); final OperatorTable table = new OperatorTable(FUNCTIONS()); final LogicalPlanPersistence lp = dbContext.getLpPersistence(); final CatalogService registry = dbContext.getCatalogService(); final QueryContext context = Mockito.mock(QueryContext.class); when(context.getSession()).thenReturn(userSession); when(context.getLpPersistence()).thenReturn(lp); when(context.getCatalogService()).thenReturn(registry); when(context.getFunctionRegistry()).thenReturn(FUNCTIONS()); when(context.getSession()).thenReturn(UserSession.Builder.newBuilder().setSupportComplexTypes(true).build()); when(context.getCurrentEndpoint()).thenReturn(NodeEndpoint.getDefaultInstance()); when(context.getActiveEndpoints()).thenReturn(ImmutableList.of(NodeEndpoint.getDefaultInstance())); when(context.getPlannerSettings()).thenReturn(new PlannerSettings(dbContext.getConfig(), queryOptions, dbContext.getClusterResourceInformation())); when(context.getOptions()).thenReturn(queryOptions); when(context.getConfig()).thenReturn(DEFAULT_SABOT_CONFIG); when(context.getOperatorTable()).thenReturn(table); when(context.getAllocator()).thenReturn(allocator); when(context.getExecutionControls()).thenReturn(executionControls); when(context.getMaterializationProvider()).thenReturn(Mockito.mock(MaterializationDescriptorProvider.class)); return context; }
private static CompleteType getLiteralSchema(QueryContext context, byte[] bytes) { try( BufferAllocator allocator = context.getAllocator().newChildAllocator("convert-from-json-sampling", 0, 1024*1024); BufferManager bufferManager = new BufferManagerImpl(allocator); ArrowBuf data = allocator.buffer(bytes.length); VectorContainer container = new VectorContainer(allocator); VectorAccessibleComplexWriter vc = new VectorAccessibleComplexWriter(container) ){ data.writeBytes(bytes); JsonReader jsonReader = new JsonReader(bufferManager.getManagedBuffer(), false, false, false); jsonReader.setSource(bytes); ComplexWriter writer = new ComplexWriterImpl("dummy", vc); writer.setPosition(0); ReadState state = jsonReader.write(writer); if(state == ReadState.END_OF_STREAM){ throw new EOFException("Unexpected arrival at end of JSON literal stream"); } container.buildSchema(); return CompleteType.fromField(container.getSchema().getFields().get(0)); }catch(Exception ex){ throw UserException.validationError(ex).message("Failure while trying to parse JSON literal.").build(logger); } } }
private OperatorContextImpl createContext(Writer writer) { BufferAllocator allocator = context.getAllocator().newChildAllocator("direct-command", 0, Long.MAX_VALUE); final OperatorStats stats = new OperatorStats(new OpProfileDef(0,0,0), allocator); final OperatorContextImpl oc = new OperatorContextImpl( context.getConfig(), FragmentHandle.newBuilder().setQueryId(context.getQueryId()).setMajorFragmentId(0).setMinorFragmentId(0).build(), writer, allocator, allocator, null, stats, null, null, context.getFunctionRegistry(), null, context.getOptions(), context.getNamespaceService(), null, NodeDebugContextProvider.NOOP, 60000, null, ImmutableList.of()); return oc; } }
@Override public Object execute() throws Exception { observer.execStarted(null); final BatchSchema schema = PojoRecordReader.getSchema(handler.getResultType()); final CollectingOutcomeListener listener = new CollectingOutcomeListener(); try(BufferAllocator allocator = context.getAllocator().newChildAllocator("direct-command", 0, Long.MAX_VALUE); VectorContainer vc = VectorContainer.create(allocator, schema); BufferManager manager = new BufferManagerImpl(allocator); final PojoRecordReader<T> reader = new PojoRecordReader<>(handler.getResultType(), result.iterator()); ) { reader.setup(new VectorContainerMutator(vc, manager)); int count = 0; while( (count = reader.next()) != 0){ vc.setRecordCount(count); final WritableBatch w = WritableBatch.get(vc); QueryData header = QueryData.newBuilder() // .setQueryId(context.getQueryId()) // .setRowCount(count) // .setDef(w.getDef()).build(); QueryWritableBatch batch = new QueryWritableBatch(header, w.getBuffers()); listener.increment(); observer.execDataArrived(listener, batch); } listener.waitForFinish(); } return null; }
try( OperatorContextImpl oc = ocx; BufferAllocator allocator = context.getAllocator().newChildAllocator("direct-writer", 0, Long.MAX_VALUE); VectorContainer vc = VectorContainer.create(allocator, schema); BufferManager manager = new BufferManagerImpl(allocator);