如何让Scala脚本在CSV/JSON Schema中推断自定义列类型？

Absolutely feasible! Spark gives you tons of flexibility to extend its built-in schema inference logic for custom data types—here's a practical approach to implement this in Scala, no pre-defined case classes required:

Core Idea #

Spark’s default CSV schema inference handles basic types (String, Integer, Double, DateTime) out of the box. We’ll build on this by:

1. First letting Spark infer the base schema
2. Sampling each column’s data to check against custom regex rules
3. Tagging columns with custom type metadata (or even mapping to custom UDTs if you need deeper integration)

Step-by-Step Implementation #

1. Load CSV with Base Schema Inference

Start by loading your CSV file using Spark’s default inference to get the basic types:

import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.types._

val spark = SparkSession.builder()
  .appName("CustomSchemaInference")
  .master("local[*]") // Remove this line for cluster deployments
  .getOrCreate()

// Load CSV with default schema inference (adjust sampling ratio for large files)
val rawDF = spark.read
  .option("header", "true")
  .option("inferSchema", "true")
  .option("samplingRatio", "0.1")
  .csv("path/to/your/csv/files")

2. Define Custom Type Regex Rules

Map your custom type names to their corresponding regex patterns. Tweak these regexes to match your exact data formats:

// Custom type regex definitions (adjust based on your data's actual format)
val customTypeRules = Map(
  "SSN" -> """^\d{3}-\d{2}-\d{4}$""".r, // US SSN format
  "VIN-ID" -> """^[A-HJ-NPR-Z0-9]{17}$""".r, // Standard 17-character vehicle VIN
  "PhoneNumber" -> """^\+?\d{1,3}[-.\s]?\(?\d{3}\)?[-.\s]?\d{3}[-.\s]?\d{4}$""".r, // Flexible international phone format
  "ID" -> """(?=(.*[A-Za-z].*){5,})(?=(.*\d.*){2,})""".r // At least 5 letters + 2 digits
)

3. Build a Custom Type Detector

Create a helper function to check if a column’s data matches a custom regex. We’ll sample data to avoid full-table scans for efficiency:

// Check if a column's sampled data matches a custom regex rule
def matchesCustomType(sampleData: Iterable[String], regex: scala.util.matching.Regex): Boolean = {
  val sample = sampleData.take(100) // Use first 100 rows for quick validation
  // Adjust to partial match (e.g., 90%+ of samples) if your data has minor inconsistencies
  sample.nonEmpty && sample.forall(regex.matches(_))
}

4. Update Schema with Custom Type Metadata

Loop through each column, check against your rules, and add custom type tags to the schema’s metadata:

// Get sampled data for each column (convert to String for regex matching)
val columnSamples = rawDF.columns.map(colName => {
  // Use sampling for large datasets to avoid pulling all data to the driver
  val sampledData = rawDF.select(colName)
    .sample(withReplacement = false, fraction = 0.05)
    .as[String]
    .collect()
    .toIterable
  (colName, sampledData)
}).toMap

// Build updated schema with custom type metadata
val enhancedSchema = StructType(rawDF.schema.fields.map(field => {
  val colName = field.name
  val sampleData = columnSamples(colName)

  // Only check String columns for custom types (since other base types are already inferred)
  val customType = if (field.dataType == StringType) {
    customTypeRules.find { case (typeName, regex) =>
      matchesCustomType(sampleData, regex)
    }.map(_._1)
  } else {
    None
  }

  // Add custom type tag to the field's metadata
  val updatedMetadata = customType match {
    case Some(ct) => new MetadataBuilder()
      .putString("customType", ct)
      .build()
    case None => field.metadata
  }

  field.copy(metadata = updatedMetadata)
}))

// Create a new DataFrame with the enhanced schema
val dfWithCustomTypes = spark.createDataFrame(rawDF.rdd, enhancedSchema)

5. Use the Custom Type Tags

You can now access the custom type metadata to handle columns differently in your pipeline:

// Example: Filter and list all columns marked as "ID" type
val idColumns = dfWithCustomTypes.schema.fields
  .filter(_.metadata.contains("customType"))
  .filter(_.metadata.getString("customType") == "ID")
  .map(_.name)

println(s"Detected ID columns: ${idColumns.mkString(", ")}")

Key Notes #

• Sampling Efficiency: For large datasets, stick to sample() instead of collect() to avoid overwhelming the driver node. Adjust the sample fraction based on your data size and accuracy needs.
• Regex Flexibility: Modify the regex patterns to match your actual data’s quirks (e.g., allow optional hyphens in SSNs, or regional phone number variations).
• Custom UDTs (Optional): If you need deeper integration (like custom serialization/deserialization), you can create a UserDefinedType for each custom type and map columns to it instead of just adding metadata. This is more complex but useful for advanced use cases.
• Dirty Data Handling: If your data has inconsistent formats, adjust the matchesCustomType function to use a threshold (e.g., 90% of samples match) instead of strict forall validation.

内容的提问来源于stack exchange，提问作者CodeHunter