[SPARK-43105][CONNECT] Abbreviate Bytes and Strings in proto message

[zhengruifeng]

What changes were proposed in this pull request?

To abbreviate the BYTES and STRING fields in proto message.

Note that the repeated and map<...> fields are always skipped for now

Why are the changes needed?

1, for abbreviation:

In [6]: spark.createDataFrame(range(0, 1000)).show()

In [7]: query = "SELECT /* " + "bla" * 8192 + " */ 1"

In [8]: spark.sql(query).show()

before:

after:

2, Message.toString may cause OOM when the message is large
This PR try to abbreviate the bytes and string, which are the main parts of LocalRelation and PythonUDF

Does this PR introduce any user-facing change?

yes, when BYTES and STRING fields are too long, abbreviate them and show the size

How was this patch tested?

manually check

[HyukjinKwon]

I think the method name should be like abbreviateBytes

[zhengruifeng]

we can also change other fields in this method via pattern matching, but right now it only abbreviate the bytes fields, so abbreviateBytes LGTM

[zhengruifeng]

I removed this try-catch for test purpose (A PyTorch UT failed due to OOM before), will add it back to be more conservative

[HyukjinKwon]

I think it's fine to remove it (since you added it to address this specific case before?). I don't mind removing it

[zhengruifeng]

yes, I added it for that PyTorch test case, in which the size of UDF is 47mb and cause OOM

but I am not very sure whether there are some other unknown edge cases that can also cause failure, so I personally prefer adding the try-catch back before merge it.

[zhengruifeng]

also cc @grundprinzip

[grundprinzip]

This creates at least one additional copy of the string that we might be able to reduce by passing the bytestring directly into the createByteString method?

[zhengruifeng]

The ByteString doesn't provide a slicing or view method, so I think we have to copy.
But we just copy a few (8 here) bytes, so should be fine

[grundprinzip]

I'm confused about this logic, assume a short string with size < NUM_FIRST_BYTES this goes into the else branch and now createByteString will return ********(redacted, size=23) is this really expected? Shouldn't this just show the short string instead?

[zhengruifeng]

it is to avoid showing all the raw data in LocalRelation

[zhengruifeng]

ok, will just show the original short string in this case.

[grundprinzip]

I really appreciate the change, but I think there might be a bug in the logic.

[zhengruifeng]

ok, on second thought, I think we should narrow this PR to abbreviation only.

I think we can support redaction as followings in the future:

{
...
      case (field: FieldDescriptor, relation: proto.LocalRelation)
          if field.getJavaType == FieldDescriptor.JavaType.MESSAGE && relation != null =>
        builder.setField(field, redactLocalRelation(relation))

      case (field: FieldDescriptor, msg: Message)
          if field.getJavaType == FieldDescriptor.JavaType.MESSAGE && msg != null =>
      ...
...
}

private def redactLocalRelation(relation: proto.LocalRelation): proto.LocalRelation = {

....
}

[grundprinzip]

since this no longer "redacted", what about making the format string something like:

"[truncated(size=XXX)]"

[HyukjinKwon]

Yeah, let's just focus on abbreviating instead of redacting. This code path would likely have to change before 4.0 for better UI in any event.

[zhengruifeng]

done

[grundprinzip]

Similar here

"$prefix[truncated(size=XXX)]"

[zhengruifeng]

done

[HyukjinKwon]

Merged to master.