Java tutorial
/** * Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.com) * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package com.linkedin.pinot.core.data.readers; import java.io.File; import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Random; import java.util.concurrent.TimeUnit; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; import org.testng.Assert; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; import org.testng.annotations.Test; import com.google.common.io.Files; import com.linkedin.pinot.common.data.DimensionFieldSpec; import com.linkedin.pinot.common.data.FieldSpec; import com.linkedin.pinot.common.data.FieldSpec.DataType; import com.linkedin.pinot.common.data.MetricFieldSpec; import com.linkedin.pinot.common.data.Schema; import com.linkedin.pinot.common.data.TimeFieldSpec; import com.linkedin.pinot.common.data.TimeGranularitySpec; import com.linkedin.pinot.core.data.GenericRow; import com.linkedin.pinot.core.indexsegment.generator.SegmentGeneratorConfig; import com.linkedin.pinot.core.segment.creator.impl.SegmentIndexCreationDriverImpl; /** * Tests the PinotSegmentRecordReader to check that the records being generated * are the same as the records used to create the segment */ public class PinotSegmentRecordReaderTest { private String segmentName; private Schema schema; private String segmentOutputDir; private File segmentIndexDir; private List<GenericRow> rows; private TestRecordReader recordReader; private static String D_SV_1 = "d_sv_1"; private static String D_MV_1 = "d_mv_1"; private static String M1 = "m1"; private static String M2 = "m2"; private static String TIME = "t"; @BeforeClass public void setup() throws Exception { segmentName = "pinotSegmentRecordReaderTest"; schema = createPinotSchema(); segmentOutputDir = Files.createTempDir().toString(); segmentIndexDir = new File(segmentOutputDir, segmentName); rows = createTestData(); recordReader = new TestRecordReader(rows, schema); createSegment(); } private List<GenericRow> createTestData() { List<GenericRow> rows = new ArrayList<>(); Random random = new Random(); Map<String, Object> fields; for (int i = 0; i < 10000; i++) { fields = new HashMap<>(); fields.put(D_SV_1, D_SV_1 + "_" + RandomStringUtils.randomAlphabetic(2)); Object[] d2Array = new Object[5]; for (int j = 0; j < 5; j++) { d2Array[j] = D_MV_1 + "_" + j + "_" + RandomStringUtils.randomAlphabetic(2); } fields.put(D_MV_1, d2Array); fields.put(M1, Math.abs(random.nextInt())); fields.put(M2, Math.abs(random.nextFloat())); fields.put(TIME, Math.abs(random.nextLong())); GenericRow row = new GenericRow(); row.init(fields); rows.add(row); } return rows; } private Schema createPinotSchema() { Schema testSchema = new Schema(); testSchema.setSchemaName("schema"); FieldSpec spec; spec = new DimensionFieldSpec(D_SV_1, DataType.STRING, true); testSchema.addField(spec); spec = new DimensionFieldSpec(D_MV_1, DataType.STRING, false); testSchema.addField(spec); spec = new MetricFieldSpec(M1, DataType.INT); testSchema.addField(spec); spec = new MetricFieldSpec(M2, DataType.FLOAT); testSchema.addField(spec); spec = new TimeFieldSpec(new TimeGranularitySpec(DataType.LONG, TimeUnit.HOURS, TIME)); testSchema.addField(spec); return testSchema; } private void createSegment() throws Exception { SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(schema); segmentGeneratorConfig.setTableName(segmentName); segmentGeneratorConfig.setOutDir(segmentOutputDir); segmentGeneratorConfig.setSegmentName(segmentName); SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); driver.init(segmentGeneratorConfig, recordReader); driver.build(); if (!segmentIndexDir.exists()) { throw new IllegalStateException("Segment generation failed"); } } @Test public void testPinotSegmentRecordReader() throws Exception { List<GenericRow> outputRows = new ArrayList<>(); PinotSegmentRecordReader pinotSegmentRecordReader = new PinotSegmentRecordReader(segmentIndexDir); pinotSegmentRecordReader.init(); while (pinotSegmentRecordReader.hasNext()) { outputRows.add(pinotSegmentRecordReader.next()); } pinotSegmentRecordReader.close(); Assert.assertEquals(outputRows.size(), rows.size(), "Number of rows returned by PinotSegmentRecordReader is incorrect"); for (int i = 0; i < outputRows.size(); i++) { GenericRow outputRow = outputRows.get(i); GenericRow row = rows.get(i); Assert.assertEquals(outputRow.getValue(D_SV_1), row.getValue(D_SV_1)); Assert.assertEquals(outputRow.getValue(D_MV_1), row.getValue(D_MV_1)); Assert.assertEquals(outputRow.getValue(M1), row.getValue(M1)); Assert.assertEquals(outputRow.getValue(M2), row.getValue(M2)); Assert.assertEquals(outputRow.getValue(TIME), row.getValue(TIME)); } } @AfterClass public void cleanup() { FileUtils.deleteQuietly(new File(segmentOutputDir)); } }