fix: properly set partition column (#407)

## Summary

- #381 introduced the ability
to configure a partition column at the node-level. This PR simply fixes
a missed spot on the plumbing of the new StagingQuery attribute.

## Checklist
- [ ] Added Unit Tests
- [ ] Covered by existing CI
- [ ] Integration tested
- [ ] Documentation update
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Enhanced the query builder to support specifying a partition column,
providing greater customization for query formation and partitioning.
- **Improvements**
- Improved handling of partition columns by introducing a fallback
mechanism to ensure valid values are used when necessary.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

<!-- av pr metadata
This information is embedded by the av CLI when creating PRs to track
the status of stacks when using Aviator. Please do not delete or edit
this section of the PR.
```
{"parent":"main","parentHead":"","trunk":"main"}
```
-->

---------

Co-authored-by: Thomas Chow <[email protected]>

Author: tchow-zlai

api/src/main/scala/ai/chronon/api/Builders.scala

@@ -307,12 +307,14 @@ object Builders {
         query: String = null,
         metaData: MetaData = null,
         startPartition: String = null,
-        setups: Seq[String] = null
+        setups: Seq[String] = null,
+        partitionColumn: String = null
     ): StagingQuery = {
       val stagingQuery = new StagingQuery()
       stagingQuery.setQuery(query)
       stagingQuery.setMetaData(metaData)
       stagingQuery.setStartPartition(startPartition)
+      stagingQuery.setPartitionColumn(partitionColumn)
       if (setups != null) stagingQuery.setSetups(setups.toJava)
       stagingQuery
     }

spark/src/main/scala/ai/chronon/spark/StagingQuery.scala

@@ -35,11 +35,12 @@ class StagingQuery(stagingQueryConf: api.StagingQuery, endPartition: String, tab
     .map(_.toScala.toMap)
     .orNull
 
-  private val partitionCols: Seq[String] = Seq(tableUtils.partitionColumn) ++
-    Option(stagingQueryConf.metaData.customJsonLookUp(key = "additional_partition_cols"))
-      .getOrElse(new java.util.ArrayList[String]())
-      .asInstanceOf[java.util.ArrayList[String]]
-      .toScala
+  private val partitionCols: Seq[String] =
+    Seq(Option(stagingQueryConf.getPartitionColumn).getOrElse(tableUtils.partitionColumn)) ++
+      Option(stagingQueryConf.metaData.customJsonLookUp(key = "additional_partition_cols"))
+        .getOrElse(new java.util.ArrayList[String]())
+        .asInstanceOf[java.util.ArrayList[String]]
+        .toScala
 
   def computeStagingQuery(stepDays: Option[Int] = None,
                           enableAutoExpand: Option[Boolean] = Some(true),