Java Code Examples for org.apache.spark.sql.types.DataTypes#TimestampType

private static DataType convertType(org.apache.kylin.metadata.datatype.DataType type) {
    if (type.isTimeFamily())
        return DataTypes.TimestampType;

    if (type.isDateTimeFamily())
        return DataTypes.DateType;

    if (type.isIntegerFamily())
        return DataTypes.LongType;

    if (type.isNumberFamily())
        return DataTypes.createDecimalType(19, 4);

    if (type.isStringFamily())
        return DataTypes.StringType;

    if (type.isBoolean())
        return DataTypes.BooleanType;

    throw new IllegalArgumentException("KAP data type: " + type + " can not be converted to spark's type.");
}
 

private DataType mapDataType(List<StructField> datasetFields, String column, String typeConfig) {

        DataType currentDatatype = getCurrentDataType(datasetFields, column);

        // when typeConfig is null (no config for this column), return the current DataType
        if(typeConfig == null) {
            return currentDatatype;
        }

        switch (typeConfig) {
            case "integer":
                return DataTypes.IntegerType;
            case "long":
                return DataTypes.LongType;
            case "double":
                return DataTypes.DoubleType;
            case "boolean":
                return DataTypes.BooleanType;
            case "date":
                return DataTypes.DateType;
            case "timestamp":
                return DataTypes.TimestampType;
            default:
                return DataTypes.StringType;
        }
    }
 

@Override
public Row call(final SimpleFeature feature) throws Exception {
  final Object[] fields = new Serializable[schema.size()];

  for (int i = 0; i < schema.size(); i++) {
    final Object fieldObj = feature.getAttribute(i);
    if (fieldObj != null) {
      final StructField structField = schema.apply(i);
      if (structField.name().equals("geom")) {
        fields[i] = fieldObj;
      } else if (structField.dataType() == DataTypes.TimestampType) {
        fields[i] = new Timestamp(((Date) fieldObj).getTime());
      } else if (structField.dataType() != null) {
        fields[i] = fieldObj;
      } else {
        LOGGER.error("Unexpected attribute in field(" + structField.name() + "): " + fieldObj);
      }
    }
  }

  return new GenericRowWithSchema(fields, schema);
}
 

private static AttributeDescriptor attrDescFromStructField(
    final AttributeTypeBuilder attrBuilder,
    final StructField field) {
  if (field.name().equals("geom")) {
    return attrBuilder.binding(Geometry.class).nillable(false).buildDescriptor("geom");
  }
  if (field.dataType() == DataTypes.StringType) {
    return attrBuilder.binding(String.class).buildDescriptor(field.name());
  } else if (field.dataType() == DataTypes.DoubleType) {
    return attrBuilder.binding(Double.class).buildDescriptor(field.name());
  } else if (field.dataType() == DataTypes.FloatType) {
    return attrBuilder.binding(Float.class).buildDescriptor(field.name());
  } else if (field.dataType() == DataTypes.LongType) {
    return attrBuilder.binding(Long.class).buildDescriptor(field.name());
  } else if (field.dataType() == DataTypes.IntegerType) {
    return attrBuilder.binding(Integer.class).buildDescriptor(field.name());
  } else if (field.dataType() == DataTypes.BooleanType) {
    return attrBuilder.binding(Boolean.class).buildDescriptor(field.name());
  } else if (field.dataType() == DataTypes.TimestampType) {
    return attrBuilder.binding(Date.class).buildDescriptor(field.name());
  }

  return null;
}
 

public static List<StructField> indexrSchemaToSparkSchema(SegmentSchema schema) {
    List<StructField> fields = new ArrayList<>();
    for (ColumnSchema cs : schema.getColumns()) {
        DataType dataType;
        switch (cs.getSqlType()) {
            case INT:
                dataType = DataTypes.IntegerType;
                break;
            case BIGINT:
                dataType = DataTypes.LongType;
                break;
            case FLOAT:
                dataType = DataTypes.FloatType;
                break;
            case DOUBLE:
                dataType = DataTypes.DoubleType;
                break;
            case VARCHAR:
                dataType = DataTypes.StringType;
                break;
            case DATE:
                dataType = DataTypes.DateType;
                break;
            case DATETIME:
                dataType = DataTypes.TimestampType;
                break;
            default:
                throw new IllegalStateException("Unsupported type: " + cs.getSqlType());
        }
        fields.add(new StructField(cs.getName(), dataType, scala.Boolean.box(false), Metadata.empty()));
    }
    return fields;
}
 

@Test
public void testToRowValueTimestamp() {
  DataType field = DataTypes.TimestampType;

  DateTime dateObj = DateTime.parse("2017-01-01T00:00:00"); // Pass-thru the TZ
  Timestamp sqlTimestamp = new Timestamp(dateObj.getMillis());

  assertEquals("Invalid Long", sqlTimestamp, RowUtils.toRowValue(dateObj.getMillis(), field));
  assertEquals("Invalid String", sqlTimestamp, RowUtils.toRowValue("2017-001", field)); // ISO Date format
  assertEquals("Invalid Date", sqlTimestamp, RowUtils.toRowValue(dateObj.toDate(), field));
  assertEquals("Invalid DateTime", sqlTimestamp, RowUtils.toRowValue(dateObj, field));

  // Test custom timestamp format parsing
  Map<RowUtils.RowValueMetadata, Object> metadataNull = Maps.newHashMap();
  Map<RowUtils.RowValueMetadata, Object> metadataEmpty = Maps.newHashMap();
  Map<RowUtils.RowValueMetadata, Object> metadataFormat = Maps.newHashMap();
  Set<String> empty = Sets.newHashSet();
  Set<String> formats = Sets.newHashSet();
  formats.add("yyyy-MM-dd HH:mm:ss.SSSSS");
  metadataNull.put(RowUtils.RowValueMetadata.TIMESTAMP_FORMATS, null);
  metadataEmpty.put(RowUtils.RowValueMetadata.TIMESTAMP_FORMATS, empty);
  metadataFormat.put(RowUtils.RowValueMetadata.TIMESTAMP_FORMATS, formats);
  assertEquals("Invalid null metadata", sqlTimestamp, RowUtils.toRowValue("2017-01-01T00:00:00", field, null));
  assertEquals("Invalid null format set in metadata", sqlTimestamp, RowUtils.toRowValue("2017-01-01T00:00:00", field, metadataNull));
  assertEquals("Invalid format set", sqlTimestamp, RowUtils.toRowValue("2017-01-01 00:00:00.00000", field, metadataFormat));
  assertEquals("Invalid empty format set", sqlTimestamp, RowUtils.toRowValue("2017-01-01T00:00:00", field, metadataEmpty));

  thrown.expect(RuntimeException.class);
  thrown.expectMessage(CoreMatchers.containsString("Invalid or unrecognized input format"));
  RowUtils.toRowValue(123, field);
}
 

private static SimpleFeatureDataType attrDescToDataType(final AttributeDescriptor attrDesc) {
  boolean isGeom = false;
  DataType dataTypeOut = DataTypes.NullType;

  if (attrDesc.getType().getBinding().equals(String.class)) {

    dataTypeOut = DataTypes.StringType;
  } else if (attrDesc.getType().getBinding().equals(Double.class)) {
    dataTypeOut = DataTypes.DoubleType;
  } else if (attrDesc.getType().getBinding().equals(Float.class)) {
    dataTypeOut = DataTypes.FloatType;
  } else if (attrDesc.getType().getBinding().equals(Long.class)) {
    dataTypeOut = DataTypes.LongType;
  } else if (attrDesc.getType().getBinding().equals(Integer.class)) {
    dataTypeOut = DataTypes.IntegerType;
  } else if (attrDesc.getType().getBinding().equals(Boolean.class)) {
    dataTypeOut = DataTypes.BooleanType;
  } else if (attrDesc.getType().getBinding().equals(Date.class)) {
    dataTypeOut = DataTypes.TimestampType;
  }

  // Custom geometry types get WKB encoding
  else if (Geometry.class.isAssignableFrom(attrDesc.getType().getBinding())) {
    dataTypeOut = GeoWaveSpatialEncoders.geometryUDT;
    isGeom = true;
  }

  return new SimpleFeatureDataType(dataTypeOut, isGeom);
}
 

private DataType sparkFromArrow(FieldType fieldType) {
  switch (fieldType.getType().getTypeID()) {
    case Null:
      return DataTypes.NullType;
    case Struct:
      throw new UnsupportedOperationException("have not implemented Struct type yet");
    case List:
      throw new UnsupportedOperationException("have not implemented List type yet");
    case FixedSizeList:
      throw new UnsupportedOperationException("have not implemented FixedSizeList type yet");
    case Union:
      throw new UnsupportedOperationException("have not implemented Union type yet");
    case Int:
      ArrowType.Int intType = (ArrowType.Int) fieldType.getType();
      int bitWidth = intType.getBitWidth();
      if (bitWidth == 8) {
        return DataTypes.ByteType;
      } else if (bitWidth == 16) {
        return DataTypes.ShortType;
      } else if (bitWidth == 32) {
        return DataTypes.IntegerType;
      } else if (bitWidth == 64) {
        return DataTypes.LongType;
      }
      throw new UnsupportedOperationException("unknown int type with bitwidth " + bitWidth);
    case FloatingPoint:
      ArrowType.FloatingPoint floatType = (ArrowType.FloatingPoint) fieldType.getType();
      FloatingPointPrecision precision = floatType.getPrecision();
      switch (precision) {
        case HALF:
        case SINGLE:
          return DataTypes.FloatType;
        case DOUBLE:
          return DataTypes.DoubleType;
      }
    case Utf8:
      return DataTypes.StringType;
    case Binary:
    case FixedSizeBinary:
      return DataTypes.BinaryType;
    case Bool:
      return DataTypes.BooleanType;
    case Decimal:
      throw new UnsupportedOperationException("have not implemented Decimal type yet");
    case Date:
      return DataTypes.DateType;
    case Time:
      return DataTypes.TimestampType; //note i don't know what this will do!
    case Timestamp:
      return DataTypes.TimestampType;
    case Interval:
      return DataTypes.CalendarIntervalType;
    case NONE:
      return DataTypes.NullType;
  }
  throw new IllegalStateException("Unexpected value: " + fieldType);
}
 

static DataType getSparkType(Type type)
{
    if (type instanceof ParameterizedType && ((ParameterizedType) type).getRawType() == Map.class) {
        Type[] arguments = ((ParameterizedType) type).getActualTypeArguments();

        return DataTypes.createMapType(getSparkType(arguments[0]), getSparkType(arguments[1]));
    }
    else if (type instanceof ParameterizedType && ((ParameterizedType) type).getRawType() == List.class) {
        DataType dataType = getSparkType(((ParameterizedType) type).getActualTypeArguments()[0]);

        return DataTypes.createArrayType(dataType);
    }
    else {
        if (type == String.class) {
            return DataTypes.StringType;
        }
        else if (type == int.class || type == Integer.class) {
            return DataTypes.IntegerType;
        }
        else if (type == long.class || type == Long.class) {
            return DataTypes.LongType;
        }
        else if (type == boolean.class || type == Boolean.class) {
            return DataTypes.BooleanType;
        }
        else if (type == double.class || type == Double.class) {
            return DataTypes.DoubleType;
        }
        else if (type == float.class || type == Float.class) {
            return DataTypes.FloatType;
        }
        else if (type == byte.class || type == Byte.class) {
            return DataTypes.ByteType;
        }
        else if (type == Timestamp.class) {
            return DataTypes.TimestampType;
        }
        else if (type == Date.class) {
            return DataTypes.DateType;
        }
        else if (type == byte[].class || type == Byte[].class) {
            return DataTypes.BinaryType;
        }
        else {
            throw new IllegalArgumentException("this TYPE " + type + " have't support!");
        }
    }
}
 

public DataType getResultType(Class<?>[] signature)
{
    return Arrays.equals(signature, new Class[] {String.class, String.class}) ? DataTypes.TimestampType : DataTypes.StringType;
}
 

public DataType getResultType(Class<?>[] signature)
{
    return signature.length == 2 ? DataTypes.StringType : DataTypes.TimestampType;
}
 

private static DataType getType(String type) {
    DataType dataType = DataTypes.NullType;
    switch (type.toLowerCase()) {
        case "string":
            dataType = DataTypes.StringType;
            break;
        case "integer":
            dataType = DataTypes.IntegerType;
            break;
        case "long":
            dataType = DataTypes.LongType;
            break;
        case "double":
            dataType = DataTypes.DoubleType;
            break;
        case "float":
            dataType = DataTypes.FloatType;
            break;
        case "short":
            dataType = DataTypes.ShortType;
            break;
        case "date":
            dataType = DataTypes.DateType;
            break;
        case "timestamp":
            dataType = DataTypes.TimestampType;
            break;
        case "boolean":
            dataType = DataTypes.BooleanType;
            break;
        case "binary":
            dataType = DataTypes.BinaryType;
            break;
        case "byte":
            dataType = DataTypes.ByteType;
            break;
        default:
            throw new ConfigRuntimeException("Throw data type exception, unknown type: " + type);
    }
    return dataType;
}
 

private StructType schemaFor(KuduTable table) {
  List<StructField> fields = Lists.newArrayList();

  for (ColumnSchema columnSchema : table.getSchema().getColumns()) {
    DataType fieldType;

    switch (columnSchema.getType()) {
      case DOUBLE:
        fieldType = DataTypes.DoubleType;
        break;
      case FLOAT:
        fieldType = DataTypes.FloatType;
        break;
      case INT8:
        fieldType = DataTypes.ByteType;
        break;
      case INT16:
        fieldType = DataTypes.ShortType;
        break;
      case INT32:
        fieldType = DataTypes.IntegerType;
        break;
      case INT64:
        fieldType = DataTypes.LongType;
        break;
      case STRING:
        fieldType = DataTypes.StringType;
        break;
      case BOOL:
        fieldType = DataTypes.BooleanType;
        break;
      case BINARY:
        fieldType = DataTypes.BinaryType;
        break;
      case UNIXTIME_MICROS:
        fieldType = DataTypes.TimestampType;
        break;
      case DECIMAL:
        int precision = columnSchema.getTypeAttributes().getPrecision();
        int scale = columnSchema.getTypeAttributes().getScale();
        fieldType = DataTypes.createDecimalType(precision, scale);
        break;
      default:
        throw new RuntimeException("Unsupported Kudu column type: " + columnSchema.getType());
    }

    fields.add(DataTypes.createStructField(columnSchema.getName(), fieldType, true));
  }

  return DataTypes.createStructType(fields);
}
 

/**
 * Convert Avro Types into their associated DataType.
 *
 * @param schemaType Avro Schema.Type
 * @return DataType representation
 */
public static DataType dataTypeFor(Schema schemaType) {
  LOG.trace("Converting Schema[{}] to DataType", schemaType);

  // Unwrap "optional" unions to the base type
  boolean isOptional = isNullable(schemaType);

  if (isOptional) {
    // if only 2 items in the union, then "unwrap," otherwise, it's a full union and should be rendered as such
    if (schemaType.getTypes().size() == 2) {
      LOG.trace("Unwrapping simple 'optional' union for {}", schemaType);
      for (Schema s : schemaType.getTypes()) {
        if (s.getType().equals(NULL)) {
          continue;
        }
        // Unwrap
        schemaType = s;
        break;
      }
    }
  }

  // Convert supported LogicalTypes
  if (null != schemaType.getLogicalType()) {
    LogicalType logicalType = schemaType.getLogicalType();
    switch (logicalType.getName()) {
      case "date" :
        return DataTypes.DateType;
      case "timestamp-millis" :
        return DataTypes.TimestampType;
      case "decimal" :
        LogicalTypes.Decimal decimal = (LogicalTypes.Decimal) logicalType;
        return DataTypes.createDecimalType(decimal.getPrecision(), decimal.getScale());
      default:
        // Pass-thru
        LOG.warn("Unsupported LogicalType[{}], continuing with underlying base type", logicalType.getName());
    }
  }

  switch (schemaType.getType()) {
    case RECORD:
      // StructType
      List<StructField> structFieldList = Lists.newArrayListWithCapacity(schemaType.getFields().size());
      for (Field f : schemaType.getFields()) {
        structFieldList.add(DataTypes.createStructField(f.name(), dataTypeFor(f.schema()), isNullable(f.schema())));
      }
      return DataTypes.createStructType(structFieldList);
    case ARRAY:
      Schema elementType = schemaType.getElementType();
      return DataTypes.createArrayType(dataTypeFor(elementType), isNullable(elementType));
    case MAP:
      Schema valueType = schemaType.getValueType();
      return DataTypes.createMapType(DataTypes.StringType, dataTypeFor(valueType), isNullable(valueType));
    case UNION:
      // StructType of members
      List<StructField> unionFieldList = Lists.newArrayListWithCapacity(schemaType.getTypes().size());
      int m = 0;
      for (Schema u : schemaType.getTypes()) {
        unionFieldList.add(DataTypes.createStructField("member" + m++, dataTypeFor(u), isNullable(u)));
      }
      return DataTypes.createStructType(unionFieldList);
    case FIXED:
    case BYTES:
      return DataTypes.BinaryType;
    case ENUM:
    case STRING:
      return DataTypes.StringType;
    case INT:
      return DataTypes.IntegerType;
    case LONG:
      return DataTypes.LongType;
    case FLOAT:
      return DataTypes.FloatType;
    case DOUBLE:
      return DataTypes.DoubleType;
    case BOOLEAN:
      return DataTypes.BooleanType;
    case NULL:
      return DataTypes.NullType;
    default:
      throw new RuntimeException(String.format("Unrecognized or unsupported Avro Type conversion: %s", schemaType));
  }
}
 

public static DataType getSparkDataType(String typeString) {
  DataType type;

  String prec_scale_regex_groups = "\\s*(decimal)\\s*\\(\\s*(\\d+)\\s*,\\s*(\\d+)\\s*\\)\\s*";
  Pattern prec_scale_regex_pattern = Pattern.compile(prec_scale_regex_groups);
  Matcher prec_scale_regex_matcher = prec_scale_regex_pattern.matcher(typeString);

  if (prec_scale_regex_matcher.matches()) {
    int precision = Integer.parseInt(prec_scale_regex_matcher.group(2)); 
    int scale = Integer.parseInt(prec_scale_regex_matcher.group(3)); 
    type = DataTypes.createDecimalType(precision, scale);
  }
  else {
    switch (typeString) {
      case DECIMAL:
        type = DataTypes.createDecimalType();
        break;
      case STRING:
        type = DataTypes.StringType;
        break;
      case FLOAT:
        type = DataTypes.FloatType;
        break;
      case DOUBLE:
        type = DataTypes.DoubleType;
        break;
      case BYTE:
        type = DataTypes.ByteType;
        break;
      case SHORT:
        type = DataTypes.ShortType;
        break;
      case INT:
        type = DataTypes.IntegerType;
        break;
      case LONG:
        type = DataTypes.LongType;
        break;
      case BOOLEAN:
        type = DataTypes.BooleanType;
        break;
      case BINARY:
        type = DataTypes.BinaryType;
        break;
      case DATE:
        type = DataTypes.DateType;
        break;
      case TIMESTAMP:
        type = DataTypes.TimestampType;
        break;
      default:
        throw new RuntimeException("Unsupported or unrecognized field type: " + typeString);
    } 
  }

  return type;
}
 

public void writeResults(String typeName) {
  if (typeName == null) {
    typeName = DEFAULT_TYPE_NAME;
    LOGGER.warn(
        "Using default type name (adapter id): '" + DEFAULT_TYPE_NAME + "' for SQL output");
  }

  final StructType schema = results.schema();
  final SimpleFeatureType featureType = SchemaConverter.schemaToFeatureType(schema, typeName);

  final SimpleFeatureBuilder sfBuilder = new SimpleFeatureBuilder(featureType);

  final FeatureDataAdapter featureAdapter = new FeatureDataAdapter(featureType);

  final DataStore featureStore = outputDataStore.createDataStore();
  final Index featureIndex =
      new SpatialDimensionalityTypeProvider().createIndex(new SpatialOptions());
  featureStore.addType(featureAdapter, featureIndex);
  try (Writer writer = featureStore.createWriter(featureAdapter.getTypeName())) {

    final List<Row> rows = results.collectAsList();

    for (int r = 0; r < rows.size(); r++) {
      final Row row = rows.get(r);

      for (int i = 0; i < schema.fields().length; i++) {
        final StructField field = schema.apply(i);
        final Object rowObj = row.apply(i);
        if (rowObj != null) {
          if (field.name().equals("geom")) {
            final Geometry geom = (Geometry) rowObj;

            sfBuilder.set("geom", geom);
          } else if (field.dataType() == DataTypes.TimestampType) {
            final long millis = ((Timestamp) rowObj).getTime();
            final Date date = new Date(millis);

            sfBuilder.set(field.name(), date);
          } else {
            sfBuilder.set(field.name(), rowObj);
          }
        }
      }

      final SimpleFeature sf = sfBuilder.buildFeature("result-" + nf.format(r));

      writer.write(sf);
    }
  }
}