Java tutorial
/** * Licensed to the Apache Software Foundation (ASF) under one or more * contributor license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright ownership. * The ASF licenses this file to You 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 org.apache.solr.response; import org.apache.commons.csv.CSVPrinter; import org.apache.commons.csv.CSVStrategy; import org.apache.lucene.document.Document; import org.apache.lucene.document.Fieldable; import org.apache.solr.common.SolrDocument; import org.apache.solr.common.SolrDocumentList; import org.apache.solr.common.SolrException; import org.apache.solr.common.params.SolrParams; import org.apache.solr.common.util.DateUtil; import org.apache.solr.common.util.FastWriter; import org.apache.solr.common.util.NamedList; import org.apache.solr.request.SolrQueryRequest; import org.apache.solr.request.uninverted.UnInvertedField; import org.apache.solr.schema.FieldType; import org.apache.solr.schema.SchemaField; import org.apache.solr.schema.StrField; import org.apache.solr.search.DocIterator; import org.apache.solr.search.DocList; import org.apache.solr.search.SolrIndexSearcher; import java.io.CharArrayWriter; import java.io.IOException; import java.io.Writer; import java.util.*; /** * @version $Id: CSVResponseWriter.java 979838 2010-07-27 20:13:15Z yonik $ */ public class CSVResponseWriter implements QueryResponseWriter { public void init(NamedList n) { } public void write(Writer writer, SolrQueryRequest req, SolrQueryResponse rsp) throws IOException { CSVWriter w = new CSVWriter(writer, req, rsp); try { w.writeResponse(); } finally { w.close(); } } public String getContentType(SolrQueryRequest request, SolrQueryResponse response) { // using the text/plain allows this to be viewed in the browser easily return CONTENT_TYPE_TEXT_UTF8; } } class CSVWriter extends TextResponseWriter { static String SEPARATOR = "separator"; static String ENCAPSULATOR = "encapsulator"; static String ESCAPE = "escape"; static String CSV = "csv."; static String CSV_SEPARATOR = CSV + SEPARATOR; static String CSV_ENCAPSULATOR = CSV + ENCAPSULATOR; static String CSV_ESCAPE = CSV + ESCAPE; static String MV = CSV + "mv."; static String MV_SEPARATOR = MV + SEPARATOR; static String MV_ENCAPSULATOR = MV + ENCAPSULATOR; static String MV_ESCAPE = MV + ESCAPE; static String CSV_NULL = CSV + "null"; static String CSV_HEADER = CSV + "header"; static String CSV_NEWLINE = CSV + "newline"; char[] sharedCSVBuf = new char[8192]; // prevent each instance from creating it's own buffer class CSVSharedBufPrinter extends CSVPrinter { public CSVSharedBufPrinter(Writer out, CSVStrategy strategy) { super(out, strategy); super.buf = sharedCSVBuf; } public void reset() { super.newLine = true; // update our shared buf in case a new bigger one was allocated sharedCSVBuf = super.buf; } } // allows access to internal buf w/o copying it static class OpenCharArrayWriter extends CharArrayWriter { public char[] getInternalBuf() { return buf; } } // Writes all data to a char array, // allows access to internal buffer, and allows fast resetting. static class ResettableFastWriter extends FastWriter { OpenCharArrayWriter cw = new OpenCharArrayWriter(); char[] result; int resultLen; public ResettableFastWriter() { super(new OpenCharArrayWriter()); cw = (OpenCharArrayWriter) sink; } public void reset() { cw.reset(); pos = 0; } public void freeze() throws IOException { if (cw.size() > 0) { flush(); result = cw.getInternalBuf(); resultLen = cw.size(); } else { result = buf; resultLen = pos; } } public int getFrozenSize() { return resultLen; } public char[] getFrozenBuf() { return result; } } static class CSVField { String name; SchemaField sf; CSVSharedBufPrinter mvPrinter; // printer used to encode multiple values in a single CSV value // used to collect values List<Fieldable> values = new ArrayList<Fieldable>(1); // low starting amount in case there are many fields int tmp; } int pass; Map<String, CSVField> csvFields = new LinkedHashMap<String, CSVField>(); Calendar cal; // for formatting date objects CSVStrategy strategy; // strategy for encoding the fields of documents CSVPrinter printer; ResettableFastWriter mvWriter = new ResettableFastWriter(); // writer used for multi-valued fields String NullValue; boolean returnScore = false; public CSVWriter(Writer writer, SolrQueryRequest req, SolrQueryResponse rsp) { super(writer, req, rsp); } public void writeResponse() throws IOException { SolrParams params = req.getParams(); strategy = new CSVStrategy(',', '"', CSVStrategy.COMMENTS_DISABLED, CSVStrategy.ESCAPE_DISABLED, false, false, false, true); CSVStrategy strat = strategy; String sep = params.get(CSV_SEPARATOR); if (sep != null) { if (sep.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid separator:'" + sep + "'"); strat.setDelimiter(sep.charAt(0)); } String nl = params.get(CSV_NEWLINE); if (nl != null) { if (nl.length() == 0) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid newline:'" + nl + "'"); strat.setPrinterNewline(nl); } String encapsulator = params.get(CSV_ENCAPSULATOR); String escape = params.get(CSV_ESCAPE); if (encapsulator != null) { if (encapsulator.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid encapsulator:'" + encapsulator + "'"); strat.setEncapsulator(encapsulator.charAt(0)); } if (escape != null) { if (escape.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid escape:'" + escape + "'"); strat.setEscape(escape.charAt(0)); if (encapsulator == null) { strat.setEncapsulator(CSVStrategy.ENCAPSULATOR_DISABLED); } } if (strat.getEscape() == '\\') { // If the escape is the standard backslash, then also enable // unicode escapes (it's harmless since 'u' would not otherwise // be escaped. strat.setUnicodeEscapeInterpretation(true); } printer = new CSVPrinter(writer, strategy); CSVStrategy mvStrategy = new CSVStrategy(strategy.getDelimiter(), CSVStrategy.ENCAPSULATOR_DISABLED, CSVStrategy.COMMENTS_DISABLED, '\\', false, false, false, false); strat = mvStrategy; sep = params.get(MV_SEPARATOR); if (sep != null) { if (sep.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid mv separator:'" + sep + "'"); strat.setDelimiter(sep.charAt(0)); } encapsulator = params.get(MV_ENCAPSULATOR); escape = params.get(MV_ESCAPE); if (encapsulator != null) { if (encapsulator.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid mv encapsulator:'" + encapsulator + "'"); strat.setEncapsulator(encapsulator.charAt(0)); if (escape == null) { strat.setEscape(CSVStrategy.ESCAPE_DISABLED); } } escape = params.get(MV_ESCAPE); if (escape != null) { if (escape.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid mv escape:'" + escape + "'"); strat.setEscape(escape.charAt(0)); // encapsulator will already be disabled if it wasn't specified } returnScore = returnFields != null && returnFields.contains("score"); boolean needListOfFields = returnFields == null || returnFields.size() == 0 || (returnFields.size() == 1 && returnScore) || returnFields.contains("*"); Collection<String> fields = returnFields; Object responseObj = rsp.getValues().get("response"); if (needListOfFields) { if (responseObj instanceof SolrDocumentList) { // get the list of fields from the SolrDocumentList fields = new LinkedHashSet<String>(); for (SolrDocument sdoc : (SolrDocumentList) responseObj) { fields.addAll(sdoc.getFieldNames()); } } else { // get the list of fields from the index fields = req.getSearcher().getFieldNames(); } if (returnScore) { fields.add("score"); } else { fields.remove("score"); } } CSVSharedBufPrinter csvPrinterMV = new CSVSharedBufPrinter(mvWriter, mvStrategy); for (String field : fields) { if (field.equals("score")) { CSVField csvField = new CSVField(); csvField.name = "score"; csvFields.put("score", csvField); continue; } SchemaField sf = schema.getFieldOrNull(field); if (sf == null) { FieldType ft = new StrField(); sf = new SchemaField(field, ft); } // if we got the list of fields from the index, only list stored fields if (returnFields == null && sf != null && !sf.stored()) { continue; } // check for per-field overrides sep = params.get("f." + field + '.' + CSV_SEPARATOR); encapsulator = params.get("f." + field + '.' + CSV_ENCAPSULATOR); escape = params.get("f." + field + '.' + CSV_ESCAPE); CSVSharedBufPrinter csvPrinter = csvPrinterMV; if (sep != null || encapsulator != null || escape != null) { // create a new strategy + printer if there were any per-field overrides strat = (CSVStrategy) mvStrategy.clone(); if (sep != null) { if (sep.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid mv separator:'" + sep + "'"); strat.setDelimiter(sep.charAt(0)); } if (encapsulator != null) { if (encapsulator.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid mv encapsulator:'" + encapsulator + "'"); strat.setEncapsulator(encapsulator.charAt(0)); if (escape == null) { strat.setEscape(CSVStrategy.ESCAPE_DISABLED); } } if (escape != null) { if (escape.length() != 1) throw new SolrException(SolrException.ErrorCode.BAD_REQUEST, "Invalid mv escape:'" + escape + "'"); strat.setEscape(escape.charAt(0)); if (encapsulator == null) { strat.setEncapsulator(CSVStrategy.ENCAPSULATOR_DISABLED); } } csvPrinter = new CSVSharedBufPrinter(mvWriter, strat); } CSVField csvField = new CSVField(); csvField.name = field; csvField.sf = sf; csvField.mvPrinter = csvPrinter; csvFields.put(field, csvField); } NullValue = params.get(CSV_NULL, ""); if (params.getBool(CSV_HEADER, true)) { for (CSVField csvField : csvFields.values()) { printer.print(csvField.name); } printer.println(); } if (responseObj instanceof DocList) { writeDocList(null, (DocList) responseObj, null, null); } else if (responseObj instanceof SolrDocumentList) { writeSolrDocumentList(null, (SolrDocumentList) responseObj, null, null); } } @Override public void close() throws IOException { if (printer != null) printer.flush(); super.close(); } @Override public void writeNamedList(String name, NamedList val) throws IOException { } @Override public void writeDoc(String name, Document doc, Set<String> returnFields, float score, boolean includeScore) throws IOException { pass++; for (Fieldable field : doc.getFields()) { CSVField csvField = csvFields.get(field.name()); if (csvField == null) continue; if (csvField.tmp != pass) { csvField.tmp = pass; csvField.values.clear(); } csvField.values.add(field); } for (CSVField csvField : csvFields.values()) { if (csvField.name.equals("score")) { writeFloat("score", score); continue; } if (csvField.tmp != pass) { writeNull(csvField.name); continue; } if (csvField.sf.multiValued() || csvField.values.size() > 1) { mvWriter.reset(); csvField.mvPrinter.reset(); // switch the printer to use the multi-valued one CSVPrinter tmp = printer; printer = csvField.mvPrinter; for (Fieldable fval : csvField.values) { csvField.sf.getType().write(this, csvField.name, fval); } printer = tmp; // restore the original printer mvWriter.freeze(); printer.print(mvWriter.getFrozenBuf(), 0, mvWriter.getFrozenSize(), true); } else { assert csvField.values.size() == 1; csvField.sf.getType().write(this, csvField.name, csvField.values.get(0)); } } printer.println(); } //NOTE: a document cannot currently contain another document List tmpList; @Override public void writeSolrDocument(String name, SolrDocument doc, Set<String> returnFields, Map pseudoFields) throws IOException { if (tmpList == null) { tmpList = new ArrayList(1); tmpList.add(null); } for (CSVField csvField : csvFields.values()) { Object val = doc.getFieldValue(csvField.name); int nVals = val instanceof Collection ? ((Collection) val).size() : (val == null ? 0 : 1); if (nVals == 0) { writeNull(csvField.name); continue; } if ((csvField.sf != null && csvField.sf.multiValued()) || nVals > 1) { Collection values; // normalize to a collection if (val instanceof Collection) { values = (Collection) val; } else { tmpList.set(0, val); values = tmpList; } mvWriter.reset(); csvField.mvPrinter.reset(); // switch the printer to use the multi-valued one CSVPrinter tmp = printer; printer = csvField.mvPrinter; for (Object fval : values) { writeVal(csvField.name, fval); } printer = tmp; // restore the original printer mvWriter.freeze(); printer.print(mvWriter.getFrozenBuf(), 0, mvWriter.getFrozenSize(), true); } else { // normalize to first value if (val instanceof Collection) { Collection values = (Collection) val; val = values.iterator().next(); } writeVal(csvField.name, val); } } printer.println(); } @Override public void writeDocList(String name, DocList ids, Set<String> fields, Map otherFields) throws IOException { int sz = ids.size(); SolrIndexSearcher searcher = req.getSearcher(); DocIterator iterator = ids.iterator(); for (int i = 0; i < sz; i++) { int id = iterator.nextDoc(); Document doc = searcher.doc(id, fields); writeDoc(null, doc, fields, (returnScore ? iterator.score() : 0.0f), returnScore); } } Map scoreMap = new HashMap(1); @Override public void writeSolrDocumentList(String name, SolrDocumentList docs, Set<String> fields, Map otherFields) throws IOException { for (SolrDocument doc : docs) { writeSolrDocument(name, doc, fields, otherFields); } } @Override public void writeStr(String name, String val, boolean needsEscaping) throws IOException { printer.print(val, needsEscaping); } @Override public void writeMap(String name, Map val, boolean excludeOuter, boolean isFirstVal) throws IOException { } @Override public void writeArray(String name, Object[] val) throws IOException { } @Override public void writeArray(String name, Iterator val) throws IOException { } @Override public void writeNull(String name) throws IOException { printer.print(NullValue); } @Override public void writeInt(String name, String val) throws IOException { printer.print(val, false); } @Override public void writeLong(String name, String val) throws IOException { printer.print(val, false); } @Override public void writeBool(String name, String val) throws IOException { printer.print(val, false); } @Override public void writeFloat(String name, String val) throws IOException { printer.print(val, false); } @Override public void writeDouble(String name, String val) throws IOException { printer.print(val, false); } @Override public void writeDate(String name, Date val) throws IOException { StringBuilder sb = new StringBuilder(25); cal = DateUtil.formatDate(val, cal, sb); writeDate(name, sb.toString()); } @Override public void writeDate(String name, String val) throws IOException { printer.print(val, false); } @Override public void writeShort(String name, String val) throws IOException { printer.print(val, false); } @Override public void writeByte(String name, String val) throws IOException { printer.print(val, false); } }