Search in sources :

Example 1 with FieldProjector

use of org.ojai.util.FieldProjector in project drill by apache.

the class MaprDBJsonRecordReader method transformColumns.

@Override
protected Collection<SchemaPath> transformColumns(Collection<SchemaPath> columns) {
    Set<SchemaPath> transformed = Sets.newLinkedHashSet();
    Set<SchemaPath> encodedSchemaPathSet = Sets.newLinkedHashSet();
    if (disablePushdown) {
        transformed.add(SchemaPath.STAR_COLUMN);
        includeId = true;
    } else {
        if (isStarQuery()) {
            transformed.add(SchemaPath.STAR_COLUMN);
            includeId = true;
            if (isSkipQuery() && !disableCountOptimization) {
                // `SELECT COUNT(*)` query
                idOnly = true;
                scannedFields = ID_ONLY_PROJECTION;
            }
        } else {
            Set<FieldPath> scannedFieldsSet = Sets.newTreeSet();
            Set<FieldPath> projectedFieldsSet = null;
            for (SchemaPath column : columns) {
                if (EncodedSchemaPathSet.isEncodedSchemaPath(column)) {
                    encodedSchemaPathSet.add(column);
                } else {
                    transformed.add(column);
                    if (!DOCUMENT_SCHEMA_PATH.equals(column)) {
                        FieldPath fp = getFieldPathForProjection(column);
                        scannedFieldsSet.add(fp);
                    } else {
                        projectWholeDocument = true;
                    }
                }
            }
            if (projectWholeDocument) {
                // we do not want to project the fields from the encoded field path list
                // hence make a copy of the scannedFieldsSet here for projection.
                projectedFieldsSet = new ImmutableSet.Builder<FieldPath>().addAll(scannedFieldsSet).build();
            }
            if (encodedSchemaPathSet.size() > 0) {
                Collection<SchemaPath> decodedSchemaPaths = EncodedSchemaPathSet.decode(encodedSchemaPathSet);
                // add them to scanned set or clear the scanned set if all fields were requested.
                for (SchemaPath column : decodedSchemaPaths) {
                    if (column.equals(SchemaPath.STAR_COLUMN)) {
                        includeId = true;
                        scannedFieldsSet.clear();
                        break;
                    }
                    scannedFieldsSet.add(getFieldPathForProjection(column));
                }
            }
            if (scannedFieldsSet.size() > 0) {
                if (includesIdField(scannedFieldsSet)) {
                    includeId = true;
                }
                scannedFields = scannedFieldsSet.toArray(new FieldPath[scannedFieldsSet.size()]);
            }
            if (disableCountOptimization) {
                idOnly = (scannedFields == null);
            }
            if (projectWholeDocument) {
                projector = new FieldProjector(projectedFieldsSet);
            }
        }
    }
    return transformed;
}
Also used : FieldProjector(org.ojai.util.FieldProjector) ImmutableSet(org.apache.drill.shaded.guava.com.google.common.collect.ImmutableSet) SchemaPath(org.apache.drill.common.expression.SchemaPath) FieldPath(org.ojai.FieldPath)

Aggregations

SchemaPath (org.apache.drill.common.expression.SchemaPath)1 ImmutableSet (org.apache.drill.shaded.guava.com.google.common.collect.ImmutableSet)1 FieldPath (org.ojai.FieldPath)1 FieldProjector (org.ojai.util.FieldProjector)1