Search in sources :

Example 1 with Dimension

use of io.tiledb.java.api.Dimension in project TileDB-Spark by TileDB-Inc.

the class TileDBBatchWrite method writeArraySchema.

private static void writeArraySchema(Context ctx, URI uri, StructType sparkSchema, TileDBDataSourceOptions options) throws TileDBError {
    ArrayType type = ArrayType.TILEDB_SPARSE;
    try (ArraySchema arraySchema = new ArraySchema(ctx, type);
        Domain domain = new Domain(ctx)) {
        String[] dimNames = TileDBWriteSchema.getSchemaDimensionOptions(sparkSchema, options);
        StructField[] sparkFields = sparkSchema.fields();
        for (int dimIdx = 0; dimIdx < dimNames.length; dimIdx++) {
            String dimName = dimNames[dimIdx];
            int idx = sparkSchema.fieldIndex(dimName);
            try (Dimension dim = TileDBWriteSchema.toDimension(ctx, dimName, dimIdx, sparkFields[idx].dataType(), options)) {
                domain.addDimension(dim);
            }
        }
        // set domain
        arraySchema.setDomain(domain);
        // add attributes
        for (StructField field : sparkFields) {
            // skip over dims
            if (Arrays.stream(dimNames).anyMatch(field.name()::equals)) {
                continue;
            }
            try (Attribute attr = TileDBWriteSchema.toAttribute(ctx, field, options)) {
                arraySchema.addAttribute(attr);
            }
        }
        // set schema tile / layouts and remaining attributes
        Optional<Layout> schemaCellOrder = options.getSchemaCellOrder();
        if (schemaCellOrder.isPresent()) {
            arraySchema.setCellOrder(schemaCellOrder.get());
        }
        Optional<Layout> schemaTileOrder = options.getSchemaTileOrder();
        if (schemaTileOrder.isPresent()) {
            arraySchema.setTileOrder(schemaTileOrder.get());
        }
        // schema filters
        Optional<List<Pair<String, Integer>>> coordsFilters = options.getSchemaCoordsFilterList();
        if (coordsFilters.isPresent()) {
            try (FilterList filterList = TileDBWriteSchema.createTileDBFilterList(ctx, coordsFilters.get())) {
                arraySchema.setCoodsFilterList(filterList);
            }
        }
        Optional<List<Pair<String, Integer>>> offsetsFilters = options.getSchemaOffsetsFilterList();
        if (coordsFilters.isPresent()) {
            try (FilterList filterList = TileDBWriteSchema.createTileDBFilterList(ctx, offsetsFilters.get())) {
                arraySchema.setOffsetsFilterList(filterList);
            }
        }
        // set capacity
        Optional<Long> schemaCapacity = options.getSchemaCapacity();
        if (schemaCapacity.isPresent()) {
            arraySchema.setCapacity(schemaCapacity.get());
        }
        // set allows dups
        if (options.getSchemaAllowDups().isPresent() && options.getSchemaAllowDups().get())
            arraySchema.setAllowDups(1);
        arraySchema.check();
        Array.create(uri.toString(), arraySchema);
    }
}
Also used : Attribute(io.tiledb.java.api.Attribute) FilterList(io.tiledb.java.api.FilterList) Dimension(io.tiledb.java.api.Dimension) ArrayType(io.tiledb.java.api.ArrayType) ArraySchema(io.tiledb.java.api.ArraySchema) StructField(org.apache.spark.sql.types.StructField) Layout(io.tiledb.java.api.Layout) FilterList(io.tiledb.java.api.FilterList) List(java.util.List) Domain(io.tiledb.java.api.Domain)

Aggregations

ArraySchema (io.tiledb.java.api.ArraySchema)1 ArrayType (io.tiledb.java.api.ArrayType)1 Attribute (io.tiledb.java.api.Attribute)1 Dimension (io.tiledb.java.api.Dimension)1 Domain (io.tiledb.java.api.Domain)1 FilterList (io.tiledb.java.api.FilterList)1 Layout (io.tiledb.java.api.Layout)1 List (java.util.List)1 StructField (org.apache.spark.sql.types.StructField)1