public JobDataFragmentWrapper(int offsetInJobResults, com.dremio.service.jobs.JobDataFragment delegate) { this.offsetInJobResults = offsetInJobResults; this.delegate = delegate; this.nameToColumns = getColumnsFromSchema(delegate.getSchema()); this.columns = ImmutableList.copyOf(nameToColumns.values()); }
@Override public void serialize(JobData jobData, JsonGenerator generator, SerializerProvider serializerProvider) throws IOException, JsonProcessingException { generator.writeStartObject(); generator.writeFieldName("rowCount"); generator.writeNumber(jobData.getRowCount()); generator.writeFieldName("schema"); generator.writeStartArray(); for (Field field : jobData.getJobDataFragment().getSchema().getFields()) { writeField(field, generator, false); } generator.writeEndArray(); generator.writeFieldName("rows"); generator.writeStartArray(); final APIJobResultsSerializer jsonWriter = new APIJobResultsSerializer(generator); jsonWriter.setup(); for(RecordBatchHolder batchHolder : jobData.delegate.getRecordBatches()) { final EventBasedRecordWriter recordWriter = new EventBasedRecordWriter(batchHolder.getData().getContainer(), jsonWriter); for (int i = batchHolder.getStart(); i < batchHolder.getEnd(); i++) { recordWriter.writeOneRecord(i); } } generator.writeEndArray(); generator.writeEndObject(); } }
@Override public JobUI answer(InvocationOnMock invocation) throws Throwable { String query = invocation.getArgumentAt(0, SqlQuery.class).getSql(); Job job = mock(Job.class); JobData jobData = mock(JobData.class); when(job.getData()).thenReturn(jobData); if ("SELECT * FROM dataset".equals(query)) { when(jobData.getJobResultsTable()).thenReturn("jobResults.previewJob"); } else if (query.contains("jobResults.previewJob")) { if (expFilter != null) { assertTrue(query, query.contains(expFilter)); } else { assertFalse(query, query.contains("WHERE")); } JobDataFragment fragment = mock(JobDataFragment.class); when(jobData.truncate(1)).thenReturn(fragment); when(fragment.getSchema()).thenReturn( BatchSchema.newBuilder() .addField(new Field("dremio_selection_count", true, new ArrowType.Int(64, true), null)) .build() ); when(fragment.extractValue("dremio_selection_count", 0)).thenReturn(expCount); } return new JobUI(job); } }
JobDataFragment result = job1_0.getData().truncate(10); JobDataFragment storedResult = jobsService.getJob(job1_0.getJobId()).getData().truncate(10); for (Field column: result.getSchema()) { assertTrue(storedResult.getSchema().getFields().contains(column)); for(Field c : result.getSchema()) { valuesFromResult.add(result.extractValue(c.getName(), i)); for(Field c : storedResult.getSchema()) { valuesFromStored.add(storedResult.extractValue(c.getName(), j));
@Test public void testQueryOnFile() throws Exception { final JobsService jobService = l(JobsService.class); Job job = jobService.submitJob(JobRequest.newBuilder() .setSqlQuery(new SqlQuery("SELECT * FROM dachdfs_test.dir1.json.\"users.json\"", SampleDataPopulator.DEFAULT_USER_NAME)) .build(), NoOpJobStatusListener.INSTANCE); JobDataFragment jobData = job.getData().truncate(500); assertEquals(3, jobData.getReturnedRowCount()); assertEquals(2, jobData.getSchema().getFieldCount()); } }