spark-connect 0.2.0

data/proto/spark/connect/expressions.proto

@@ -0,0 +1,557 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+syntax = 'proto3';
+
+import "google/protobuf/any.proto";
+import "spark/connect/types.proto";
+import "spark/connect/common.proto";
+
+package spark.connect;
+
+option java_multiple_files = true;
+option java_package = "org.apache.spark.connect.proto";
+option go_package = "internal/generated";
+
+// Expression used to refer to fields, functions and similar. This can be used everywhere
+// expressions in SQL appear.
+message Expression {
+
+  ExpressionCommon common = 18;
+  oneof expr_type {
+    Literal literal = 1;
+    UnresolvedAttribute unresolved_attribute = 2;
+    UnresolvedFunction unresolved_function = 3;
+    ExpressionString expression_string = 4;
+    UnresolvedStar unresolved_star = 5;
+    Alias alias = 6;
+    Cast cast = 7;
+    UnresolvedRegex unresolved_regex = 8;
+    SortOrder sort_order = 9;
+    LambdaFunction lambda_function = 10;
+    Window window = 11;
+    UnresolvedExtractValue unresolved_extract_value = 12;
+    UpdateFields update_fields = 13;
+    UnresolvedNamedLambdaVariable unresolved_named_lambda_variable = 14;
+    CommonInlineUserDefinedFunction common_inline_user_defined_function = 15;
+    CallFunction call_function = 16;
+    NamedArgumentExpression named_argument_expression = 17;
+    MergeAction merge_action = 19;
+    TypedAggregateExpression typed_aggregate_expression = 20;
+    SubqueryExpression subquery_expression = 21;
+    DirectShufflePartitionID direct_shuffle_partition_id = 22;
+
+    // This field is used to mark extensions to the protocol. When plugins generate arbitrary
+    // relations they can add them here. During the planning the correct resolution is done.
+    google.protobuf.Any extension = 999;
+  }
+
+
+  // Expression for the OVER clause or WINDOW clause.
+  message Window {
+
+    // (Required) The window function.
+    Expression window_function = 1;
+
+    // (Optional) The way that input rows are partitioned.
+    repeated Expression partition_spec = 2;
+
+    // (Optional) Ordering of rows in a partition.
+    repeated SortOrder order_spec = 3;
+
+    // (Optional) Window frame in a partition.
+    //
+    // If not set, it will be treated as 'UnspecifiedFrame'.
+    WindowFrame frame_spec = 4;
+
+    // The window frame
+    message WindowFrame {
+
+      // (Required) The type of the frame.
+      FrameType frame_type = 1;
+
+      // (Required) The lower bound of the frame.
+      FrameBoundary lower = 2;
+
+      // (Required) The upper bound of the frame.
+      FrameBoundary upper = 3;
+
+      enum FrameType {
+        FRAME_TYPE_UNDEFINED = 0;
+
+        // RowFrame treats rows in a partition individually.
+        FRAME_TYPE_ROW = 1;
+
+        // RangeFrame treats rows in a partition as groups of peers.
+        // All rows having the same 'ORDER BY' ordering are considered as peers.
+        FRAME_TYPE_RANGE = 2;
+      }
+
+      message FrameBoundary {
+        oneof boundary {
+          // CURRENT ROW boundary
+          bool current_row = 1;
+
+          // UNBOUNDED boundary.
+          // For lower bound, it will be converted to 'UnboundedPreceding'.
+          // for upper bound, it will be converted to 'UnboundedFollowing'.
+          bool unbounded = 2;
+
+          // This is an expression for future proofing. We are expecting literals on the server side.
+          Expression value = 3;
+        }
+      }
+    }
+  }
+
+  // SortOrder is used to specify the  data ordering, it is normally used in Sort and Window.
+  // It is an unevaluable expression and cannot be evaluated, so can not be used in Projection.
+  message SortOrder {
+    // (Required) The expression to be sorted.
+    Expression child = 1;
+
+    // (Required) The sort direction, should be ASCENDING or DESCENDING.
+    SortDirection direction = 2;
+
+    // (Required) How to deal with NULLs, should be NULLS_FIRST or NULLS_LAST.
+    NullOrdering null_ordering = 3;
+
+    enum SortDirection {
+      SORT_DIRECTION_UNSPECIFIED = 0;
+      SORT_DIRECTION_ASCENDING = 1;
+      SORT_DIRECTION_DESCENDING = 2;
+    }
+
+    enum NullOrdering {
+      SORT_NULLS_UNSPECIFIED = 0;
+      SORT_NULLS_FIRST = 1;
+      SORT_NULLS_LAST = 2;
+    }
+  }
+
+  // Expression that takes a partition ID value and passes it through directly for use in
+  // shuffle partitioning. This is used with RepartitionByExpression to allow users to
+  // directly specify target partition IDs.
+  message DirectShufflePartitionID {
+    // (Required) The expression that evaluates to the partition ID.
+    Expression child = 1;
+  }
+
+  message Cast {
+    // (Required) the expression to be casted.
+    Expression expr = 1;
+
+    // (Required) the data type that the expr to be casted to.
+    oneof cast_to_type {
+      DataType type = 2;
+      // If this is set, Server will use Catalyst parser to parse this string to DataType.
+      string type_str = 3;
+    }
+
+    // (Optional) The expression evaluation mode.
+    EvalMode eval_mode = 4;
+
+    enum EvalMode {
+      EVAL_MODE_UNSPECIFIED = 0;
+      EVAL_MODE_LEGACY = 1;
+      EVAL_MODE_ANSI = 2;
+      EVAL_MODE_TRY = 3;
+    }
+  }
+
+  message Literal {
+    oneof literal_type {
+      DataType null = 1;
+      bytes binary = 2;
+      bool boolean = 3;
+
+      int32 byte = 4;
+      int32 short = 5;
+      int32 integer = 6;
+      int64 long = 7;
+      float float = 10;
+      double double = 11;
+      Decimal decimal = 12;
+
+      string string = 13;
+
+      // Date in units of days since the UNIX epoch.
+      int32 date = 16;
+      // Timestamp in units of microseconds since the UNIX epoch.
+      int64 timestamp = 17;
+      // Timestamp in units of microseconds since the UNIX epoch (without timezone information).
+      int64 timestamp_ntz = 18;
+
+      CalendarInterval calendar_interval = 19;
+      int32 year_month_interval = 20;
+      int64 day_time_interval = 21;
+      Array array = 22;
+      Map map = 23;
+      Struct struct = 24;
+
+      SpecializedArray specialized_array = 25;
+      Time time = 26;
+    }
+
+    // Reserved for Geometry and Geography.
+    reserved 27, 28;
+
+    // Data type information for the literal.
+    // This field is required only in the root literal message for null values or
+    // for data types (e.g., array, map, or struct) with non-trivial information.
+    // If the data_type field is not set at the root level, the data type will be
+    // inferred or retrieved from the deprecated data type fields using best efforts.
+    DataType data_type = 100;
+
+    message Decimal {
+      // the string representation.
+      string value = 1;
+      // The maximum number of digits allowed in the value.
+      // the maximum precision is 38.
+      optional int32 precision = 2;
+      // declared scale of decimal literal
+      optional int32 scale = 3;
+    }
+
+    message CalendarInterval {
+      int32 months = 1;
+      int32 days = 2;
+      int64 microseconds = 3;
+    }
+
+    message Array {
+      // (Deprecated) The element type of the array.
+      //
+      // This field is deprecated since Spark 4.1+. Use data_type field instead.
+      DataType element_type = 1 [deprecated = true];
+
+      // The literal values that make up the array elements.
+      repeated Literal elements = 2;
+    }
+
+    message Map {
+      // (Deprecated) The key type of the map.
+      //
+      // This field is deprecated since Spark 4.1+. Use data_type field instead.
+      DataType key_type = 1 [deprecated = true];
+
+      // (Deprecated) The value type of the map.
+      //
+      // This field is deprecated since Spark 4.1+ and should only be set
+      // if the data_type field is not set. Use data_type field instead.
+      DataType value_type = 2 [deprecated = true];
+
+      // The literal keys that make up the map.
+      repeated Literal keys = 3;
+
+      // The literal values that make up the map.
+      repeated Literal values = 4;
+    }
+
+    message Struct {
+      // (Deprecated) The type of the struct.
+      //
+      // This field is deprecated since Spark 4.1+ because using DataType as the type of a struct
+      // is ambiguous. Use data_type field instead.
+      DataType struct_type = 1 [deprecated = true];
+
+      // The literal values that make up the struct elements.
+      repeated Literal elements = 2;
+    }
+
+    message SpecializedArray {
+      oneof value_type {
+        Bools bools = 1;
+        Ints ints = 2;
+        Longs longs = 3;
+        Floats floats = 4;
+        Doubles doubles = 5;
+        Strings strings = 6;
+      }
+    }
+
+    message Time {
+      int64 nano = 1;
+      // The precision of this time, if omitted, uses the default value of MICROS_PRECISION.
+      optional int32 precision = 2;
+    }
+  }
+
+  // An unresolved attribute that is not explicitly bound to a specific column, but the column
+  // is resolved during analysis by name.
+  message UnresolvedAttribute {
+    // (Required) An identifier that will be parsed by Catalyst parser. This should follow the
+    // Spark SQL identifier syntax.
+    string unparsed_identifier = 1;
+
+    // (Optional) The id of corresponding connect plan.
+    optional int64 plan_id = 2;
+
+    // (Optional) The requested column is a metadata column.
+    optional bool is_metadata_column = 3;
+  }
+
+  // An unresolved function is not explicitly bound to one explicit function, but the function
+  // is resolved during analysis following Sparks name resolution rules.
+  message UnresolvedFunction {
+    // (Required) name (or unparsed name for user defined function) for the unresolved function.
+    string function_name = 1;
+
+    // (Optional) Function arguments. Empty arguments are allowed.
+    repeated Expression arguments = 2;
+
+    // (Required) Indicate if this function should be applied on distinct values.
+    bool is_distinct = 3;
+
+    // (Required) Indicate if this is a user defined function.
+    //
+    // When it is not a user defined function, Connect will use the function name directly.
+    // When it is a user defined function, Connect will parse the function name first.
+    bool is_user_defined_function = 4;
+
+    // (Optional) Indicate if this function is defined in the internal function registry.
+    // If not set, the server will try to look up the function in the internal function registry
+    // and decide appropriately.
+    optional bool is_internal = 5;
+  }
+
+  // Expression as string.
+  message ExpressionString {
+    // (Required) A SQL expression that will be parsed by Catalyst parser.
+    string expression = 1;
+  }
+
+  // UnresolvedStar is used to expand all the fields of a relation or struct.
+  message UnresolvedStar {
+
+    // (Optional) The target of the expansion.
+    //
+    // If set, it should end with '.*' and will be parsed by 'parseAttributeName'
+    // in the server side.
+    optional string unparsed_target = 1;
+
+    // (Optional) The id of corresponding connect plan.
+    optional int64 plan_id = 2;
+  }
+
+  // Represents all of the input attributes to a given relational operator, for example in
+  // "SELECT `(id)?+.+` FROM ...".
+  message UnresolvedRegex {
+    // (Required) The column name used to extract column with regex.
+    string col_name = 1;
+
+    // (Optional) The id of corresponding connect plan.
+    optional int64 plan_id = 2;
+  }
+
+  // Extracts a value or values from an Expression
+  message UnresolvedExtractValue {
+    // (Required) The expression to extract value from, can be
+    // Map, Array, Struct or array of Structs.
+    Expression child = 1;
+
+    // (Required) The expression to describe the extraction, can be
+    // key of Map, index of Array, field name of Struct.
+    Expression extraction = 2;
+  }
+
+  // Add, replace or drop a field of `StructType` expression by name.
+  message UpdateFields {
+    // (Required) The struct expression.
+    Expression struct_expression = 1;
+
+    // (Required) The field name.
+    string field_name = 2;
+
+    // (Optional) The expression to add or replace.
+    //
+    // When not set, it means this field will be dropped.
+    Expression value_expression = 3;
+  }
+
+  message Alias {
+    // (Required) The expression that alias will be added on.
+    Expression expr = 1;
+
+    // (Required) a list of name parts for the alias.
+    //
+    // Scalar columns only has one name that presents.
+    repeated string name = 2;
+
+    // (Optional) Alias metadata expressed as a JSON map.
+    optional string metadata = 3;
+  }
+
+  message LambdaFunction {
+    // (Required) The lambda function.
+    //
+    // The function body should use 'UnresolvedAttribute' as arguments, the sever side will
+    // replace 'UnresolvedAttribute' with 'UnresolvedNamedLambdaVariable'.
+    Expression function = 1;
+
+    // (Required) Function variables. Must contains 1 ~ 3 variables.
+    repeated Expression.UnresolvedNamedLambdaVariable arguments = 2;
+  }
+
+  message UnresolvedNamedLambdaVariable {
+
+    // (Required) a list of name parts for the variable. Must not be empty.
+    repeated string name_parts = 1;
+  }
+}
+
+message ExpressionCommon {
+  // (Required) Keep the information of the origin for this expression such as stacktrace.
+  Origin origin = 1;
+}
+
+message CommonInlineUserDefinedFunction {
+  // (Required) Name of the user-defined function.
+  string function_name = 1;
+  // (Optional) Indicate if the user-defined function is deterministic.
+  bool deterministic = 2;
+  // (Optional) Function arguments. Empty arguments are allowed.
+  repeated Expression arguments = 3;
+  // (Required) Indicate the function type of the user-defined function.
+  oneof function {
+    PythonUDF python_udf = 4;
+    ScalarScalaUDF scalar_scala_udf = 5;
+    JavaUDF java_udf = 6;
+  }
+  // (Required) Indicate if this function should be applied on distinct values.
+  bool is_distinct = 7;
+}
+
+message PythonUDF {
+  // (Required) Output type of the Python UDF
+  DataType output_type = 1;
+  // (Required) EvalType of the Python UDF
+  int32 eval_type = 2;
+  // (Required) The encoded commands of the Python UDF
+  bytes command = 3;
+  // (Required) Python version being used in the client.
+  string python_ver = 4;
+  // (Optional) Additional includes for the Python UDF.
+  repeated string additional_includes = 5;
+}
+
+message ScalarScalaUDF {
+  // (Required) Serialized JVM object containing UDF definition, input encoders and output encoder
+  bytes payload = 1;
+  // (Optional) Input type(s) of the UDF
+  repeated DataType inputTypes = 2;
+  // (Required) Output type of the UDF
+  DataType outputType = 3;
+  // (Required) True if the UDF can return null value
+  bool nullable = 4;
+  // (Required) Indicate if the UDF is an aggregate function
+  bool aggregate = 5;
+}
+
+message JavaUDF {
+  // (Required) Fully qualified name of Java class
+  string class_name = 1;
+
+  // (Optional) Output type of the Java UDF
+  optional DataType output_type = 2;
+
+  // (Required) Indicate if the Java user-defined function is an aggregate function
+  bool aggregate = 3;
+}
+
+message TypedAggregateExpression {
+  // (Required) The aggregate function object packed into bytes.
+  ScalarScalaUDF scalar_scala_udf = 1;
+}
+
+message CallFunction {
+  // (Required) Unparsed name of the SQL function.
+  string function_name = 1;
+
+  // (Optional) Function arguments. Empty arguments are allowed.
+  repeated Expression arguments = 2;
+}
+
+message NamedArgumentExpression {
+  // (Required) The key of the named argument.
+  string key = 1;
+
+  // (Required) The value expression of the named argument.
+  Expression value = 2;
+}
+
+message MergeAction {
+  // (Required) The action type of the merge action.
+  ActionType action_type = 1;
+
+  // (Optional) The condition expression of the merge action.
+  optional Expression condition = 2;
+
+  // (Optional) The assignments of the merge action. Required for ActionTypes INSERT and UPDATE.
+  repeated Assignment assignments = 3;
+
+  enum ActionType {
+    ACTION_TYPE_INVALID = 0;
+    ACTION_TYPE_DELETE = 1;
+    ACTION_TYPE_INSERT = 2;
+    ACTION_TYPE_INSERT_STAR = 3;
+    ACTION_TYPE_UPDATE = 4;
+    ACTION_TYPE_UPDATE_STAR = 5;
+  }
+
+  message Assignment {
+    // (Required) The key of the assignment.
+    Expression key = 1;
+
+    // (Required) The value of the assignment.
+    Expression value = 2;
+  }
+}
+
+message SubqueryExpression {
+  // (Required) The ID of the corresponding connect plan.
+  int64 plan_id = 1;
+
+  // (Required) The type of the subquery.
+  SubqueryType subquery_type = 2;
+
+  // (Optional) Options specific to table arguments.
+  optional TableArgOptions table_arg_options = 3;
+
+  // (Optional) IN subquery values.
+  repeated Expression in_subquery_values = 4;
+
+  enum SubqueryType {
+    SUBQUERY_TYPE_UNKNOWN = 0;
+    SUBQUERY_TYPE_SCALAR = 1;
+    SUBQUERY_TYPE_EXISTS = 2;
+    SUBQUERY_TYPE_TABLE_ARG = 3;
+    SUBQUERY_TYPE_IN = 4;
+  }
+
+  // Nested message for table argument options.
+  message TableArgOptions {
+    // (Optional) The way that input rows are partitioned.
+    repeated Expression partition_spec = 1;
+
+    // (Optional) Ordering of rows in a partition.
+    repeated Expression.SortOrder order_spec = 2;
+
+    // (Optional) Whether this is a single partition.
+    optional bool with_single_partition = 3;
+  }
+}

data/proto/spark/connect/ml.proto

@@ -0,0 +1,147 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+syntax = 'proto3';
+
+package spark.connect;
+
+import "spark/connect/relations.proto";
+import "spark/connect/expressions.proto";
+import "spark/connect/ml_common.proto";
+
+option java_multiple_files = true;
+option java_package = "org.apache.spark.connect.proto";
+option go_package = "internal/generated";
+
+// Command for ML
+message MlCommand {
+  oneof command {
+    Fit fit = 1;
+    Fetch fetch = 2;
+    Delete delete = 3;
+    Write write = 4;
+    Read read = 5;
+    Evaluate evaluate = 6;
+    CleanCache clean_cache = 7;
+    GetCacheInfo get_cache_info = 8;
+    CreateSummary create_summary = 9;
+    GetModelSize get_model_size = 10;
+  }
+
+  // Command for estimator.fit(dataset)
+  message Fit {
+    // (Required) Estimator information (its type should be OPERATOR_TYPE_ESTIMATOR)
+    MlOperator estimator = 1;
+    // (Optional) parameters of the Estimator
+    optional MlParams params = 2;
+    // (Required) the training dataset
+    Relation dataset = 3;
+  }
+
+  // Command to delete the cached objects which could be a model
+  // or summary evaluated by a model
+  message Delete {
+    repeated ObjectRef obj_refs = 1;
+    // if set `evict_only` to true, only evict the cached model from memory,
+    // but keep the offloaded model in Spark driver local disk.
+    optional bool evict_only = 2;
+  }
+
+  // Force to clean up all the ML cached objects
+  message CleanCache { }
+
+  // Get the information of all the ML cached objects
+  message GetCacheInfo { }
+
+  // Command to write ML operator
+  message Write {
+    // It could be an estimator/evaluator or the cached model
+    oneof type {
+      // Estimator or evaluator
+      MlOperator operator = 1;
+      // The cached model
+      ObjectRef obj_ref = 2;
+    }
+    // (Optional) The parameters of operator which could be estimator/evaluator or a cached model
+    optional MlParams params = 3;
+    // (Required) Save the ML instance to the path
+    string path = 4;
+    // (Optional) Overwrites if the output path already exists.
+    optional bool should_overwrite = 5;
+    // (Optional) The options of the writer
+    map<string, string> options = 6;
+  }
+
+  // Command to load ML operator.
+  message Read {
+    // (Required) ML operator information
+    MlOperator operator = 1;
+    // (Required) Load the ML instance from the input path
+    string path = 2;
+  }
+
+  // Command for evaluator.evaluate(dataset)
+  message Evaluate {
+    // (Required) Evaluator information (its type should be OPERATOR_TYPE_EVALUATOR)
+    MlOperator evaluator = 1;
+    // (Optional) parameters of the Evaluator
+    optional MlParams params = 2;
+    // (Required) the evaluating dataset
+    Relation dataset = 3;
+  }
+
+  // This is for re-creating the model summary when the model summary is lost
+  // (model summary is lost when the model is offloaded and then loaded back)
+  message CreateSummary {
+    ObjectRef model_ref = 1;
+    Relation dataset = 2;
+  }
+
+  // This is for query the model estimated in-memory size
+  message GetModelSize {
+    ObjectRef model_ref = 1;
+  }
+}
+
+// The result of MlCommand
+message MlCommandResult {
+  oneof result_type {
+    // The result of the attribute
+    Expression.Literal param = 1;
+    // Evaluate a Dataset in a model and return the cached ID of summary
+    string summary = 2;
+    // Operator information
+    MlOperatorInfo operator_info = 3;
+  }
+
+  // Represents an operator info
+  message MlOperatorInfo {
+    oneof type {
+      // The cached object which could be a model or summary evaluated by a model
+      ObjectRef obj_ref = 1;
+      // Operator name
+      string name = 2;
+    }
+    // (Optional) the 'uid' of a ML object
+    // Note it is different from the 'id' of a cached object.
+    optional string uid = 3;
+    // (Optional) parameters
+    optional MlParams params = 4;
+    // (Optional) warning message generated during the ML command execution
+    optional string warning_message = 5;
+  }
+}