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.tools; import com.linkedin.pinot.core.startree.StarTreeIndexNodeInterf; import com.linkedin.pinot.core.startree.StarTreeInterf; import com.linkedin.pinot.core.startree.StarTreeSerDe; import java.io.File; import java.io.FileInputStream; import java.io.FileNotFoundException; import java.util.ArrayList; import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Set; import org.apache.commons.lang3.tuple.ImmutablePair; import org.apache.commons.lang3.tuple.Pair; import org.codehaus.jackson.map.ObjectMapper; import org.codehaus.jackson.map.annotate.JsonSerialize.Inclusion; import org.restlet.Application; import org.restlet.Component; import org.restlet.Restlet; import org.restlet.data.Protocol; import org.restlet.representation.Representation; import org.restlet.representation.StringRepresentation; import org.restlet.resource.Directory; import org.restlet.resource.Get; import org.restlet.resource.ServerResource; import org.restlet.routing.Router; import org.restlet.routing.VirtualHost; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import com.google.common.collect.HashBiMap; import com.google.common.collect.MinMaxPriorityQueue; import com.linkedin.pinot.common.segment.ReadMode; import com.linkedin.pinot.core.common.Block; import com.linkedin.pinot.core.common.BlockSingleValIterator; import com.linkedin.pinot.core.common.BlockValSet; import com.linkedin.pinot.core.common.DataSource; import com.linkedin.pinot.core.indexsegment.IndexSegment; import com.linkedin.pinot.core.segment.creator.impl.V1Constants; import com.linkedin.pinot.core.segment.index.SegmentMetadataImpl; import com.linkedin.pinot.core.segment.index.loader.Loaders; import com.linkedin.pinot.core.segment.index.readers.Dictionary; public class StarTreeIndexViewer { private static final Logger LOGGER = LoggerFactory.getLogger(StarTreeIndexViewer.class); /* * MAX num children to show in the UI */ static int MAX_CHILDREN = 100; private HashBiMap<String, Integer> dimensionNameToIndexMap; private Map<String, Dictionary> dictionaries; private Map<String, BlockSingleValIterator> valueIterators; public StarTreeIndexViewer(File segmentDir) throws Exception { IndexSegment indexSegment = Loaders.IndexSegment.load(segmentDir, ReadMode.heap); dictionaries = new HashMap<String, Dictionary>(); valueIterators = new HashMap<String, BlockSingleValIterator>(); SegmentMetadataImpl metadata = new SegmentMetadataImpl(segmentDir); for (String columnName : metadata.getAllColumns()) { DataSource dataSource = indexSegment.getDataSource(columnName); dataSource.open(); Block block = dataSource.nextBlock(); BlockValSet blockValSet = block.getBlockValueSet(); BlockSingleValIterator itr = (BlockSingleValIterator) blockValSet.iterator(); valueIterators.put(columnName, itr); dictionaries.put(columnName, dataSource.getDictionary()); } File starTreeFile = new File(segmentDir, V1Constants.STAR_TREE_INDEX_FILE); StarTreeInterf tree = StarTreeSerDe.fromFile(starTreeFile, ReadMode.mmap); dimensionNameToIndexMap = tree.getDimensionNameToIndexMap(); StarTreeJsonNode jsonRoot = new StarTreeJsonNode("ROOT"); build(tree.getRoot(), jsonRoot); ObjectMapper objectMapper = new ObjectMapper(); objectMapper.getSerializationConfig().setSerializationInclusion(Inclusion.NON_NULL); String writeValueAsString = objectMapper.defaultPrettyPrintingWriter().writeValueAsString(jsonRoot); LOGGER.info(writeValueAsString); startServer(segmentDir, writeValueAsString); } private int build(StarTreeIndexNodeInterf indexNode, StarTreeJsonNode json) { Iterator<? extends StarTreeIndexNodeInterf> childrenIterator = indexNode.getChildrenIterator(); if (!childrenIterator.hasNext()) { return 0; } int childDimensionId = indexNode.getChildDimensionName(); String childDimensionName = dimensionNameToIndexMap.inverse().get(childDimensionId); Dictionary dictionary = dictionaries.get(childDimensionName); int totalChildNodes = indexNode.getNumChildren(); Comparator<Pair<String, Integer>> comparator = new Comparator<Pair<String, Integer>>() { @Override public int compare(Pair<String, Integer> o1, Pair<String, Integer> o2) { return -1 * Integer.compare(o1.getRight(), o2.getRight()); } }; MinMaxPriorityQueue<Pair<String, Integer>> queue = MinMaxPriorityQueue.orderedBy(comparator) .maximumSize(MAX_CHILDREN).create(); StarTreeJsonNode allNode = null; while (childrenIterator.hasNext()) { StarTreeIndexNodeInterf childIndexNode = childrenIterator.next(); int childDimensionValueId = childIndexNode.getDimensionValue(); String childDimensionValue = "ALL"; if (childDimensionValueId != StarTreeIndexNodeInterf.ALL) { childDimensionValue = dictionary.get(childDimensionValueId).toString(); } StarTreeJsonNode childJson = new StarTreeJsonNode(childDimensionValue); totalChildNodes += build(childIndexNode, childJson); if (childDimensionValueId != StarTreeIndexNodeInterf.ALL) { json.addChild(childJson); queue.add(ImmutablePair.of(childDimensionValue, totalChildNodes)); } else { allNode = childJson; } } //put ALL node at the end if (allNode != null) { json.addChild(allNode); } if (totalChildNodes > MAX_CHILDREN) { Iterator<Pair<String, Integer>> qIterator = queue.iterator(); Set<String> topKDimensions = new HashSet<>(); topKDimensions.add("ALL"); while (qIterator.hasNext()) { topKDimensions.add(qIterator.next().getKey()); } Iterator<StarTreeJsonNode> iterator = json.getChildren().iterator(); while (iterator.hasNext()) { StarTreeJsonNode next = iterator.next(); if (!topKDimensions.contains(next.getName())) { iterator.remove(); } } } return totalChildNodes; } public static void main(String[] args) throws Exception { if (args.length != 1) { LOGGER.error("USAGE: StarIndexViewer <segmentDirectory>"); System.exit(1); } String segmentDir = args[0]; if (!new File(segmentDir).exists()) { throw new FileNotFoundException("Missing directory:" + segmentDir); } new StarTreeIndexViewer(new File(segmentDir)); } private void startServer(final File segmentDirectory, final String json) throws Exception { Component component = new Component(); int port = 8090; component.getServers().add(Protocol.HTTP, port); component.getClients().add(Protocol.FILE); Application application = new Application() { @Override public Restlet createInboundRoot() { Router router = new Router(getContext()); StarTreeViewRestResource.json = json; router.attach("/data", StarTreeViewRestResource.class); Directory directory = new Directory(getContext(), getClass().getClassLoader().getResource("star-tree.html").toString()); router.attach(directory); return router; } }; VirtualHost defaultHost = component.getDefaultHost(); defaultHost.attach(application); component.start(); LOGGER.info("Go to http://{}:{}/ to view the star tree", VirtualHost.getLocalHostName(), port); } public static final class StarTreeViewRestResource extends ServerResource { public static String json; @Get @Override public Representation get() { return new StringRepresentation(json); } } } class StarTreeJsonNode { public StarTreeJsonNode(String name) { this.name = name; } String name; String parent; int size = 0; List<StarTreeJsonNode> children; public void addChild(StarTreeJsonNode childJsonNode) { if (children == null) { children = new ArrayList<>(); } children.add(childJsonNode); } public String getParent() { return parent; } public void setParent(String parent) { this.parent = parent; } public void addChild(String name) { if (children == null) { children = new ArrayList<>(); } children.add(new StarTreeJsonNode(name)); } public String getName() { return name; } public void setName(String name) { this.name = name; } public List<StarTreeJsonNode> getChildren() { return children; } public void setChildren(List<StarTreeJsonNode> children) { this.children = children; } }