Examples with CSVPrinter org.apache.solr.internal.csv.CSVPrinter used on opensource projects

Example 1 with CSVPrinter

use of org.apache.solr.internal.csv.CSVPrinter in project lucene-solr by apache.

the class CSVWriter method writeSolrDocument.

@Override
public void writeSolrDocument(String name, SolrDocument doc, ReturnFields returnFields, int idx) 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);
            }
            // restore the original printer
            printer = tmp;
            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();
            }
            // if field is polyfield, use the multi-valued printer to apply appropriate escaping
            if (csvField.sf != null && csvField.sf.isPolyField()) {
                mvWriter.reset();
                csvField.mvPrinter.reset();
                CSVPrinter tmp = printer;
                printer = csvField.mvPrinter;
                writeVal(csvField.name, val);
                printer = tmp;
                mvWriter.freeze();
                printer.print(mvWriter.getFrozenBuf(), 0, mvWriter.getFrozenSize(), true);
            } else {
                writeVal(csvField.name, val);
            }
        }
    }
    printer.println();
}

Example 2 with CSVPrinter

use of org.apache.solr.internal.csv.CSVPrinter in project lucene-solr by apache.

the class CSVWriter method writeResponse.

public void writeResponse() throws IOException {
    SolrParams params = req.getParams();
    strategy = new CSVStrategy(',', '"', CSVStrategy.COMMENTS_DISABLED, CSVStrategy.ESCAPE_DISABLED, false, false, false, true, "\n");
    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, "\n");
    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
    }
    Collection<String> fields = returnFields.getRequestedFieldNames();
    Object responseObj = rsp.getResponse();
    boolean returnOnlyStored = false;
    if (fields == null || returnFields.hasPatternMatching()) {
        if (responseObj instanceof SolrDocumentList) {
            // get the list of fields from the SolrDocumentList
            if (fields == null) {
                fields = new LinkedHashSet<>();
            }
            for (SolrDocument sdoc : (SolrDocumentList) responseObj) {
                fields.addAll(sdoc.getFieldNames());
            }
        } else {
            // get the list of fields from the index
            Iterable<String> all = req.getSearcher().getFieldNames();
            if (fields == null) {
                fields = Sets.newHashSet(all);
            } else {
                Iterables.addAll(fields, all);
            }
        }
        if (returnFields.wantsScore()) {
            fields.add("score");
        } else {
            fields.remove("score");
        }
        returnOnlyStored = true;
    }
    CSVSharedBufPrinter csvPrinterMV = new CSVSharedBufPrinter(mvWriter, mvStrategy);
    for (String field : fields) {
        if (!returnFields.wantsField(field)) {
            continue;
        }
        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);
        }
        // Return only stored fields, unless an explicit field list is specified
        if (returnOnlyStored && 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);
        // if polyfield and no escape is provided, add "\\" escape by default
        if (sf.isPolyField()) {
            escape = (escape == null) ? "\\" : 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 ResultContext) {
        writeDocuments(null, (ResultContext) responseObj);
    } else if (responseObj instanceof DocList) {
        ResultContext ctx = new BasicResultContext((DocList) responseObj, returnFields, null, null, req);
        writeDocuments(null, ctx);
    } else if (responseObj instanceof SolrDocumentList) {
        writeSolrDocumentList(null, (SolrDocumentList) responseObj, returnFields);
    }
}