[SPARK-47545][CONNECT] Dataset observe support for the Scala client

[xupefei]

What changes were proposed in this pull request?

This PR adds support for Dataset.observe to the Spark Connect Scala client. Note that the support here does not include listener support as it runs on the serve side.

This PR includes a small refactoring to the Observation helper class. We extracted methods that are not bound to the SparkSession to spark-api, and added two subclasses on both spark-core and spark-jvm-client.

Why are the changes needed?

Before this PR, the DF.observe method is only supported in the Python client.

Does this PR introduce any user-facing change?

Yes. The user can now issue DF.observe(name, metrics...) or DF.observe(observationObject, metrics...) to get stats of columns of a dataframe.

How was this patch tested?

Added new e2e tests.

Was this patch authored or co-authored using generative AI tooling?

Nope.

[HyukjinKwon]

Let's make sure this is matched with Python verison.

[xupefei]

Let's make sure this is matched with Python verison.

Thanks Hyukjin! Yes the behavior largely matches the Python version (internals and user-facing APIs). One difference though is the df.collectObservations() API which in Python has a different syntax df.attrs["observed_metrics"].

[ueshin]

So df.collectObservations() seems to be a new API available only in Spark Connect Scala client?

[xupefei]

So df.collectObservations() seems to be a new API available only in Spark Connect Scala client?

Yes, similar to df.attrs["observed_metrics"] which is only in the Python client.

[hvanhovell]

Should this be tied to whether or not the observation has been successfully updated? Other question under what circumstance can the metrics be empty.

[xupefei]

I had the same question when I looked at the code. In Spark Core we only de-register the Observation when some non-empty metrics are set, so I decide to keep it the same in Connect. I am not sure under which circumstance the metrics can be empty.

[xupefei]

I looked at the code. It seems it's valid to check for non-empty metrics in Spark Core:

private[spark] def onFinish(qe: QueryExecution): Unit = {
  ...
  val row: Option[Row] = qe.observedMetrics.get(name)
  val metrics: Option[Map[String, Any]] = row.map(r => r.getValuesMap[Any](r.schema.fieldNames.toImmutableArraySeq))
  if (setMetricsAndNotify(metrics)) {
    unregister()
  }
}

The option is for a case when the query finishes without metric. Is this possible?

[xupefei]

But nevertheless we don't need to handle this case in Connect because we look up using Plan Id.

[hvanhovell]

For the record. IMO the observation class should have been using a future from the get go.

[hvanhovell]

There is a bit of a twist here. So, LiteralValueProtoConverter, returns a Tuple for a nested struct. This is not really expected in a Row. We can address this in a follow-up.

[hvanhovell]

LGTM

[hvanhovell]

@xupefei there is a genuine test failure. Can you check what is going on?

[xupefei]

@xupefei there is a genuine test failure. Can you check what is going on?

It seems the test is flaky, even after the previous attempt to fix it: #45173

[xupefei]

@xupefei there is a genuine test failure. Can you check what is going on?

It seems the test is flaky, even after the previous attempt to fix it: #45173

I re-ran the test and it did pass. The remaining failure is doc generation: sh: 1: python: not found.

[hvanhovell]

Merging!