feat: TableUtils to be compatible with DataPointer (part 1) #158
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||||||||||
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| @@ -1,53 +1,22 @@ | ||||||||||||||||||
| package ai.chronon.integrations.cloud_gcp | ||||||||||||||||||
|
|
||||||||||||||||||
| import ai.chronon.spark.Format | ||||||||||||||||||
| import org.apache.spark.sql.Encoders | ||||||||||||||||||
| import org.apache.spark.sql.Row | ||||||||||||||||||
| import org.apache.spark.sql.SparkSession | ||||||||||||||||||
| import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder | ||||||||||||||||||
| import org.apache.spark.sql.execution.FileSourceScanExec | ||||||||||||||||||
| import org.apache.spark.sql.execution.datasources.PartitioningAwareFileIndex | ||||||||||||||||||
|
|
||||||||||||||||||
| case class GCS(project: String, sourceUri: String, format: String) extends Format { | ||||||||||||||||||
|
There was a problem hiding this comment. 💡 Codebase verification All GCS usages need updates Found commented-out legacy code and active usages that need to be updated with the new parameters:
🔗 Analysis chainVerify all GCS instantiations are updated Breaking change: new required parameters. Let me check the context of these usages to ensure they're properly updated. 🏁 Scripts executedThe following scripts were executed for the analysis: Script: #!/bin/bash
# Find all GCS instantiations
rg "new GCS|GCS\(" --type scala
Length of output: 576 Script: #!/bin/bash
# Get context around GCS usages in BigQueryFormat
rg "GCS\(" -A 5 -B 5 cloud_gcp/src/main/scala/ai/chronon/integrations/cloud_gcp/BigQueryFormat.scala
Length of output: 1457
tchow-zlai marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||||||||
|
|
||||||||||||||||||
| override def name: String = format | ||||||||||||||||||
|
|
||||||||||||||||||
| override def primaryPartitions(tableName: String, partitionColumn: String, subPartitionsFilter: Map[String, String])( | ||||||||||||||||||
| implicit sparkSession: SparkSession): Seq[String] = | ||||||||||||||||||
| super.primaryPartitions(tableName, partitionColumn, subPartitionsFilter) | ||||||||||||||||||
|
|
||||||||||||||||||
| override def partitions(tableName: String)(implicit sparkSession: SparkSession): Seq[Map[String, String]] = { | ||||||||||||||||||
|
|
||||||||||||||||||
| /** | ||||||||||||||||||
| * Given: | ||||||||||||||||||
|
|
@@ -70,7 +39,7 @@ case class GCS(project: String) extends Format { | |||||||||||||||||
| * | ||||||||||||||||||
| */ | ||||||||||||||||||
| val partitionSpec = sparkSession.read | ||||||||||||||||||
| .parquet(sourceUri) | ||||||||||||||||||
tchow-zlai marked this conversation as resolved.
Show resolved
Hide resolved
tchow-zlai marked this conversation as resolved.
Show resolved
Hide resolved
|
||||||||||||||||||
| .queryExecution | ||||||||||||||||||
| .sparkPlan | ||||||||||||||||||
| .asInstanceOf[FileSourceScanExec] | ||||||||||||||||||
|
|
@@ -82,16 +51,23 @@ case class GCS(project: String) extends Format { | |||||||||||||||||
| val partitionColumns = partitionSpec.partitionColumns | ||||||||||||||||||
| val partitions = partitionSpec.partitions.map(_.values) | ||||||||||||||||||
|
|
||||||||||||||||||
| val deserializer = | ||||||||||||||||||
| Encoders.row(partitionColumns).asInstanceOf[ExpressionEncoder[Row]].resolveAndBind().createDeserializer() | ||||||||||||||||||
|
|
||||||||||||||||||
|
There was a problem hiding this comment. Add error handling for schema deserialization The deserializer creation could fail if partition columns schema is invalid. - val deserializer =
- Encoders.row(partitionColumns).asInstanceOf[ExpressionEncoder[Row]].resolveAndBind().createDeserializer()
+ val deserializer = try {
+ Encoders.row(partitionColumns).asInstanceOf[ExpressionEncoder[Row]].resolveAndBind().createDeserializer()
+ } catch {
+ case e: Exception =>
+ throw new RuntimeException(s"Failed to create deserializer for partition columns: ${e.getMessage}", e)
+ }📝 Committable suggestion
Suggested change
There was a problem hiding this comment. seems like it could be helpful? There was a problem hiding this comment. @david-zlai, indeed, wrapping the deserializer creation in a try-catch block can be helpful. It ensures that if the partition columns schema is invalid or causes an exception, you’ll have a clear message indicating the failure. For example: val deserializer = try {
Encoders.row(partitionColumns).asInstanceOf[ExpressionEncoder[Row]].resolveAndBind().createDeserializer()
} catch {
case e: Exception =>
throw new RuntimeException(s"Failed to create deserializer for partition columns: ${e.getMessage}", e)
}This way, any schema-related issues will be properly surfaced and handled. Let me know if you need anything else! There was a problem hiding this comment.
|
||||||||||||||||||
| val roundTripped = sparkSession | ||||||||||||||||||
| .createDataFrame(sparkSession.sparkContext.parallelize(partitions.map(deserializer)), partitionColumns) | ||||||||||||||||||
| .collect | ||||||||||||||||||
| .toList | ||||||||||||||||||
|
|
||||||||||||||||||
| roundTripped | ||||||||||||||||||
| .map((part) => | ||||||||||||||||||
| partitionColumns.fields.toList.zipWithIndex.map { | ||||||||||||||||||
| case (field, idx) => { | ||||||||||||||||||
| val fieldName = field.name | ||||||||||||||||||
| val fieldValue = part.get(idx) | ||||||||||||||||||
| fieldName -> fieldValue.toString // Just going to cast this as a string. | ||||||||||||||||||
|
Comment on lines
+72
to
73
There was a problem hiding this comment. Handle null values in partition values Missing null check for fieldValue. - val fieldValue = part.get(idx)
- fieldName -> fieldValue.toString
+ val fieldValue = Option(part.get(idx))
+ fieldName -> fieldValue.map(_.toString).getOrElse("")📝 Committable suggestion
Suggested change
|
||||||||||||||||||
| } | ||||||||||||||||||
| }.toMap) | ||||||||||||||||||
| } | ||||||||||||||||||
|
|
||||||||||||||||||
| def createTableTypeString: String = throw new UnsupportedOperationException("GCS does not support create table") | ||||||||||||||||||
|
|
||||||||||||||||||
Uh oh!
There was an error while loading. Please reload this page.