| /* | |
| * 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.druid.spark.utils | |
| import org.apache.druid.data.input.InputRow | |
| import org.apache.druid.data.input.impl.{DimensionSchema, DimensionsSpec, DoubleDimensionSchema, | |
| FloatDimensionSchema, LongDimensionSchema, StringDimensionSchema} | |
| import org.apache.druid.java.util.common.IAE | |
| import org.apache.druid.segment.QueryableIndex | |
| import org.apache.druid.segment.column.{RowSignature, ValueType} | |
| import org.apache.druid.spark.registries.ComplexTypeRegistry | |
| import org.apache.spark.sql.catalyst.InternalRow | |
| import org.apache.spark.sql.catalyst.util.ArrayData | |
| import org.apache.spark.sql.types.{ArrayType, BinaryType, DataType, DoubleType, FloatType, | |
| IntegerType, LongType, StringType, StructField, StructType, TimestampType} | |
| import org.apache.spark.unsafe.types.UTF8String | |
| import java.util.{Collection => JCollection} | |
| import scala.collection.JavaConverters.{collectionAsScalaIterableConverter, | |
| iterableAsScalaIterableConverter, mapAsScalaMapConverter, seqAsJavaListConverter} | |
| /** | |
| * Converters and utilities for working with Spark and Druid schemas. | |
| */ | |
| object SchemaUtils { | |
| /** | |
| * Convert a COLUMNMAP representing a Druid datasource's schema as returned by | |
| * DruidMetadataClient.getClient into a Spark StructType. | |
| * | |
| * @param columnMap The Druid schema to convert into a corresponding Spark StructType. | |
| * @return The StructType equivalent of the Druid schema described by COLUMNMAP. | |
| */ | |
| def convertDruidSchemaToSparkSchema(columnMap: Map[String, (String, Boolean)]): StructType = { | |
| StructType.apply( | |
| columnMap.map { case (name, (colType, hasMultipleValues)) => | |
| val sparkType = colType match { | |
| case "LONG" => LongType | |
| case "STRING" => StringType | |
| case "DOUBLE" => DoubleType | |
| case "FLOAT" => FloatType | |
| case "TIMESTAMP" => TimestampType | |
| case complexType: String if ComplexTypeRegistry.getRegisteredMetricNames.contains(complexType) => | |
| BinaryType | |
| // Add other supported types later | |
| case _ => throw new IAE(s"Unrecognized type $colType!") | |
| } | |
| if (hasMultipleValues) { | |
| StructField(name, new ArrayType(sparkType, false)) | |
| } else { | |
| StructField(name, sparkType) | |
| } | |
| }.toSeq | |
| ) | |
| } | |
| /** | |
| * Convert a Druid INPUTROW into a Spark InternalRow with schema SCHEMA. | |
| * | |
| * @param inputRow The Druid InputRow to convert into a Spark Row for loading into a dataframe. | |
| * @param schema The schema to map INPUTROW into. | |
| * @param useDefaultNullHandling Whether to use the Druid defaults for null values or actual nulls. | |
| * @return A Spark InternalRow with schema SCHEMA and values parsed from INPUTROW. | |
| */ | |
| def convertInputRowToSparkRow( | |
| inputRow: InputRow, | |
| schema: StructType, | |
| useDefaultNullHandling: Boolean | |
| ): InternalRow = { | |
| InternalRow.fromSeq(schema.fieldNames.map { colName => | |
| if (colName == "__time") { | |
| inputRow.getTimestampFromEpoch | |
| } else { | |
| val col = inputRow.getRaw(colName) | |
| if (col != null) { | |
| schema(colName).dataType match { | |
| case _: ArrayType => | |
| val baseType = schema(colName).dataType.asInstanceOf[ArrayType].elementType | |
| col match { | |
| case collection: JCollection[_] => | |
| ArrayData.toArrayData(collection.asScala.map { elem => | |
| parseToScala(elem, baseType) | |
| }) | |
| case _ => | |
| // Single-element arrays won't be wrapped when read from Druid; need to do it here | |
| ArrayData.toArrayData(List(parseToScala(col, baseType))) | |
| } | |
| case _ => | |
| // This is slightly inefficient since some objects will already be the correct type | |
| parseToScala(col, schema(colName).dataType) | |
| } | |
| } else { | |
| if (useDefaultNullHandling) { | |
| schema(colName).dataType match { | |
| case StringType | ArrayType(StringType, _) => UTF8String.EMPTY_UTF8 | |
| case LongType | IntegerType | FloatType | DoubleType | TimestampType => 0 | |
| case _ => null // scalastyle:ignore null | |
| } | |
| } else { | |
| null // scalastyle:ignore null | |
| } | |
| } | |
| } | |
| }) | |
| } | |
| /** | |
| * Convert an object COL to the appropriate scala type for the given Spark DataType DT. | |
| * | |
| * @param col The object to convert to a suitable type. | |
| * @param dt The Spark DataType COL should be made compatible with. | |
| * @return COL parsed into a type compatible with DT. | |
| */ | |
| def parseToScala(col: Any, dt: DataType): Any = { | |
| dt match { | |
| case StringType => UTF8String.fromString(col.toString) | |
| case LongType => col match { | |
| case _: java.lang.Long | Long => col | |
| case _: String => col.asInstanceOf[String].toLong | |
| case _ => throw new IllegalArgumentException( | |
| s"Unsure how to parse ${col.getClass.toString} into a Long!" | |
| ) | |
| } | |
| case TimestampType => col // Timestamps should always come back from Druid as DateTimes | |
| case FloatType => col match { | |
| case _: java.lang.Float | Float => col | |
| case _: String => col.asInstanceOf[String].toFloat | |
| case _ => throw new IllegalArgumentException( | |
| s"Unsure how to parse ${col.getClass.toString} into a Float!" | |
| ) | |
| } | |
| case DoubleType => col match { | |
| case _: java.lang.Double | Double => col | |
| case _: String => col.asInstanceOf[String].toDouble | |
| case _ => throw new IllegalArgumentException( | |
| s"Unsure how to parse ${col.getClass.toString} into a Double!" | |
| ) | |
| } | |
| case BinaryType => | |
| if (ComplexTypeRegistry.getRegisteredSerializedClasses.contains(col.getClass)) { | |
| ComplexTypeRegistry.deserialize(col) | |
| } else { | |
| col match { | |
| case arr: Array[Byte] => | |
| arr | |
| case _ => throw new IllegalArgumentException( | |
| s"Unsure how to parse ${col.getClass.toString} into a ByteArray!" | |
| ) | |
| } | |
| } | |
| case _ => throw new IllegalArgumentException( | |
| s"$dt currently unsupported!" | |
| ) | |
| } | |
| } | |
| /** | |
| * Given a list of column names DIMENSIONS and a struct type SCHEMA, returns a list of Druid DimensionSchemas | |
| * constructed from each column named in DIMENSIONS and the properties of the corresponding field in SCHEMA. | |
| * | |
| * @param dimensions A list of column names to construct Druid DimensionSchemas for. | |
| * @param schema The Spark schema to use to determine the types of the Druid DimensionSchemas created. | |
| * @return A list of DimensionSchemas generated from the type information in SCHEMA for each dimension in DIMENSIONS. | |
| */ | |
| def convertStructTypeToDruidDimensionSchema( | |
| dimensions: Seq[String], | |
| schema: StructType | |
| ): Seq[DimensionSchema] = { | |
| schema | |
| .filter(field => dimensions.contains(field.name)) | |
| .map(field => | |
| field.dataType match { | |
| case LongType | IntegerType => new LongDimensionSchema(field.name) | |
| case FloatType => new FloatDimensionSchema(field.name) | |
| case DoubleType => new DoubleDimensionSchema(field.name) | |
| case StringType | ArrayType(StringType, _) => | |
| new StringDimensionSchema(field.name) | |
| case _ => throw new IAE( | |
| "Unsure how to create dimension from column [%s] with data type [%s]", | |
| field.name, | |
| field.dataType | |
| ) | |
| } | |
| ) | |
| } | |
| /** | |
| * Validates that the given list of DimensionSchemas align with the given Spark schema. This validation is done to | |
| * fail fast if the user-provided dimensions have different data types than the data in the source data frame to be | |
| * written. | |
| * | |
| * @param dimensions The list of DimensionSchemas to validate against SCHEMA. | |
| * @param schema The source-of-truth Spark schema to ensure compatibility with. | |
| * @throws IAE If the data types in DIMENSIONS do not align with the data types in SCHEMA. | |
| */ | |
| def validateDimensionSpecAgainstSparkSchema(dimensions: Seq[DimensionSchema], schema: StructType): Boolean = { | |
| val incompatibilities = dimensions.flatMap{dim => | |
| if (schema.fieldNames.contains(dim.getName)) { | |
| val sparkType = schema(schema.fieldIndex(dim.getName)).dataType | |
| sparkType match { | |
| case LongType | IntegerType => | |
| if (dim.getTypeName != DimensionSchema.LONG_TYPE_NAME) { | |
| Some(s"${dim.getName}: expected type ${DimensionSchema.LONG_TYPE_NAME} but was ${dim.getTypeName}!") | |
| } else { | |
| None | |
| } | |
| case FloatType => | |
| if (dim.getTypeName != DimensionSchema.FLOAT_TYPE_NAME) { | |
| Some(s"${dim.getName}: expected type ${DimensionSchema.FLOAT_TYPE_NAME} but was ${dim.getTypeName}!") | |
| } else { | |
| None | |
| } | |
| case DoubleType => | |
| if (dim.getTypeName != DimensionSchema.DOUBLE_TYPE_NAME) { | |
| Some(s"${dim.getName}: expected type ${DimensionSchema.DOUBLE_TYPE_NAME} but was ${dim.getTypeName}!") | |
| } else { | |
| None | |
| } | |
| case StringType | ArrayType(StringType, _) => | |
| if (dim.getTypeName != DimensionSchema.STRING_TYPE_NAME) { | |
| Some(s"${dim.getName}: expected type ${DimensionSchema.STRING_TYPE_NAME} but was ${dim.getTypeName}!") | |
| } else { | |
| None | |
| } | |
| } | |
| } else { | |
| None | |
| } | |
| } | |
| if (incompatibilities.nonEmpty) { | |
| throw new IAE(s"Incompatible dimensions spec provided! Offending columns: ${incompatibilities.mkString("; ")}") | |
| } | |
| // For now, the return type could just be Unit, but leaving the stubs in place for future improvement | |
| true | |
| } | |
| /** | |
| * Given a Spark schema, construct an equivalent Druid RowSignature. | |
| * | |
| * @param schema The Spark schema to generate a RowSignature from. | |
| * @return A RowSignature corresponding to SCHEMA. | |
| */ | |
| def generateRowSignatureFromSparkSchema(schema: StructType): RowSignature = { | |
| val builder = RowSignature.builder() | |
| schema.foreach{field => | |
| builder.add(field.name, getDruidValueTypeForDataType(field.dataType)) | |
| } | |
| builder.build() | |
| } | |
| /** | |
| * Return the Druid ValueType corresponding to a provided Spark DataType. | |
| * | |
| * @param dtype The Spark DataType | |
| * @return The ValueType corresponding to DTYPE. | |
| */ | |
| def getDruidValueTypeForDataType(dtype: DataType): ValueType = { | |
| dtype match { | |
| case DoubleType => ValueType.DOUBLE | |
| case FloatType => ValueType.FLOAT | |
| case LongType => ValueType.LONG | |
| case StringType => ValueType.STRING | |
| case ArrayType(DoubleType, _) => ValueType.DOUBLE_ARRAY | |
| case ArrayType(LongType, _) => ValueType.LONG_ARRAY | |
| case ArrayType(StringType, _) => ValueType.STRING_ARRAY | |
| case _ => ValueType.COMPLEX | |
| } | |
| } | |
| def getDimensionsSpecFromIndex(queryableIndex: QueryableIndex): DimensionsSpec = { | |
| val dimensionHandlers = queryableIndex.getDimensionHandlers.asScala | |
| new DimensionsSpec(queryableIndex.getAvailableDimensions.asScala.toSeq.map{ | |
| dim => | |
| val columnHolder = queryableIndex.getColumnHolder(dim) | |
| val handler = dimensionHandlers(dim) | |
| columnHolder.getCapabilities.getType match { | |
| case ValueType.DOUBLE => new DoubleDimensionSchema(dim) | |
| case ValueType.FLOAT => new FloatDimensionSchema(dim) | |
| case ValueType.LONG => new LongDimensionSchema(dim) | |
| case ValueType.STRING => new StringDimensionSchema( | |
| dim, handler.getMultivalueHandling, columnHolder.getCapabilities.hasBitmapIndexes | |
| ) | |
| case ValueType.DOUBLE_ARRAY | ValueType.LONG_ARRAY | ValueType.STRING_ARRAY | ValueType.COMPLEX => | |
| throw new IAE( | |
| s"This reader cannot process dimension [$dim] with type [${columnHolder.getCapabilities.getType}]!" | |
| ) | |
| } | |
| }.toList.asJava) | |
| } | |
| } |