spark-connect 0.3.0 → 0.4.0

data/lib/spark_connect/grouped_data.rb

@@ -11,17 +11,24 @@ module SparkConnect
   class GroupedData
     Proto = SparkConnect::Proto
 
+    NUMERIC_TYPES = [
+      Types::ByteType, Types::ShortType, Types::IntegerType, Types::LongType,
+      Types::FloatType, Types::DoubleType, Types::DecimalType,
+    ].freeze
+
     # @param df [DataFrame]
     # @param grouping [Array<Column>] grouping columns.
     # @param group_type [Symbol] a `GROUP_TYPE_*` enum symbol.
     # @param pivot_col [Column, nil]
     # @param pivot_values [Array, nil]
-    def initialize(df, grouping, group_type, pivot_col: nil, pivot_values: nil)
+    # @param grouping_sets [Array<Array<Column>>, nil] for `GROUP_TYPE_GROUPING_SETS`.
+    def initialize(df, grouping, group_type, pivot_col: nil, pivot_values: nil, grouping_sets: nil)
       @df = df
       @grouping = grouping
       @group_type = group_type
       @pivot_col = pivot_col
       @pivot_values = pivot_values
+      @grouping_sets = grouping_sets
     end
 
     # Compute aggregate expressions.
@@ -33,6 +40,8 @@ module SparkConnect
     #     `{"age" => "max", "salary" => "avg"}`.
     # @return [DataFrame]
     def agg(*exprs)
+      raise IllegalArgumentError, "agg requires at least one expression" if exprs.empty?
+
       agg_exprs =
         if exprs.size == 1 && exprs.first.is_a?(Hash)
           exprs.first.map { |col, fn| Column.invoke(fn.to_s, Functions.col(col.to_s)).to_expr }
@@ -65,22 +74,34 @@ module SparkConnect
     # @return [DataFrame]
     def min(*cols) = numeric_agg("min", cols)
 
-    # Pivot a column into multiple output columns.
+    # Pivot a column into multiple output columns. Only valid right after
+    # {DataFrame#group_by}.
     #
     # @param pivot_col [String, Column]
     # @param values [Array, nil] optional explicit pivot values (faster, deterministic).
     # @return [GroupedData]
     def pivot(pivot_col, values = nil)
-      GroupedData.new(@df, @grouping, :GROUP_TYPE_PIVOT,
-                      pivot_col: Column.to_col(pivot_col.is_a?(String) ? Functions.col(pivot_col) : pivot_col),
-                      pivot_values: values)
+      unless @group_type == :GROUP_TYPE_GROUPBY
+        what = @group_type == :GROUP_TYPE_PIVOT ? "Repeated PIVOT operation" : "PIVOT after #{group_type_name}"
+        raise NotImplementedError, "#{what} is not supported"
+      end
+      raise IllegalArgumentError, "pivot values must be an Array, got #{values.class}" unless values.nil? || values.is_a?(Array)
+
+      col = pivot_col.is_a?(Column) ? pivot_col : @df.col(pivot_col.to_s)
+      GroupedData.new(@df, @grouping, :GROUP_TYPE_PIVOT, pivot_col: col, pivot_values: values)
     end
 
     private
 
+    def group_type_name
+      @group_type.to_s.delete_prefix("GROUP_TYPE_")
+    end
+
+    # Like PySpark, aggregate every numeric column when none is named.
     def numeric_agg(fn, cols)
-      exprs = cols.flatten.map { |c| Column.invoke(fn, Functions.col(c.to_s)).to_expr }
-      build(exprs)
+      names = cols.flatten.map(&:to_s)
+      names = @df.schema.fields.select { |f| NUMERIC_TYPES.any? { |t| f.data_type.is_a?(t) } }.map(&:name) if names.empty?
+      build(names.map { |c| Column.invoke(fn, Functions.col(c)).to_expr })
     end
 
     def build(agg_exprs)
@@ -95,6 +116,11 @@ module SparkConnect
         pivot.values += @pivot_values.map { |v| Column.lit(v).to_expr.literal } if @pivot_values
         agg.pivot = pivot
       end
+      if @group_type == :GROUP_TYPE_GROUPING_SETS
+        agg.grouping_sets += (@grouping_sets || []).map do |set|
+          Proto::Aggregate::GroupingSets.new(grouping_set: set.map(&:to_expr))
+        end
+      end
       @df.build(aggregate: agg)
     end
   end

data/lib/spark_connect/merge.rb

@@ -0,0 +1,146 @@
+# frozen_string_literal: true
+
+module SparkConnect
+  # Builds and runs a `MERGE INTO` command from a source {DataFrame}. Returned by
+  # {DataFrame#merge_into}. Mirrors PySpark's `MergeIntoWriter` (Spark 4.0+).
+  #
+  # Each `when_*` clause returns a small builder whose action method
+  # (`update`, `update_all`, `delete`, `insert`, `insert_all`) records the
+  # action and returns this writer, so clauses chain:
+  #
+  # @example
+  #   source.merge_into("target", F.expr("target.id = source.id"))
+  #         .when_matched(F.expr("source.deleted")).delete
+  #         .when_matched.update("value" => F.col("source.value"))
+  #         .when_not_matched.insert_all
+  #         .when_not_matched_by_source.delete
+  #         .merge
+  class MergeIntoWriter
+    Proto = SparkConnect::Proto
+    MA = Proto::MergeAction
+
+    # @param df [DataFrame] the source.
+    # @param table [String] the target table.
+    # @param condition [Column] the merge condition.
+    def initialize(df, table, condition)
+      raise IllegalArgumentError, "merge condition must be a Column, got #{condition.class}" unless condition.is_a?(Column)
+
+      @df = df
+      @table = table.to_s
+      @condition = condition
+      @schema_evolution = false
+      @matched = []
+      @not_matched = []
+      @not_matched_by_source = []
+    end
+
+    # Actions for target rows that match a source row.
+    # @param condition [Column, nil] an additional condition for this clause.
+    # @return [WhenMatched]
+    def when_matched(condition = nil) = WhenMatched.new(self, condition)
+    alias whenMatched when_matched
+
+    # Actions for source rows without a matching target row.
+    # @param condition [Column, nil]
+    # @return [WhenNotMatched]
+    def when_not_matched(condition = nil) = WhenNotMatched.new(self, condition)
+    alias whenNotMatched when_not_matched
+
+    # Actions for target rows without a matching source row.
+    # @param condition [Column, nil]
+    # @return [WhenNotMatchedBySource]
+    def when_not_matched_by_source(condition = nil) = WhenNotMatchedBySource.new(self, condition)
+    alias whenNotMatchedBySource when_not_matched_by_source
+
+    # Let the merge evolve the target schema (add source columns).
+    # @return [self]
+    def with_schema_evolution
+      @schema_evolution = true
+      self
+    end
+    alias withSchemaEvolution with_schema_evolution
+
+    # Run the merge. @return [void]
+    def merge
+      expr = ->(action) { Proto::Expression.new(merge_action: action) }
+      cmd = Proto::MergeIntoTableCommand.new(
+        target_table_name: @table, source_table_plan: @df.relation, merge_condition: @condition.to_expr,
+        match_actions: @matched.map(&expr), not_matched_actions: @not_matched.map(&expr),
+        not_matched_by_source_actions: @not_matched_by_source.map(&expr),
+        with_schema_evolution: @schema_evolution
+      )
+      @df.session.client.execute_command(Proto::Command.new(merge_into_table_command: cmd))
+      nil
+    end
+
+    # Append a merge action of `kind` (matched, not matched, not matched by
+    # source) and return the writer.
+    # @api private
+    def add_action(kind, type, condition, assignments = nil)
+      action = MA.new(action_type: type)
+      action.condition = condition.to_expr if condition
+      if assignments
+        action.assignments += assignments.map do |key, value|
+          MA::Assignment.new(key: Functions.expr(key.to_s).to_expr, value: Column.to_col(value).to_expr)
+        end
+      end
+      instance_variable_get(:"@#{kind}") << action
+      self
+    end
+
+    # Clause builder returned by {MergeIntoWriter#when_matched}.
+    class WhenMatched
+      def initialize(writer, condition)
+        @writer = writer
+        @condition = condition
+      end
+
+      # Update every column of the matched target row from the source row.
+      # @return [MergeIntoWriter]
+      def update_all = @writer.add_action(:matched, :ACTION_TYPE_UPDATE_STAR, @condition)
+      alias updateAll update_all
+
+      # @param assignments [Hash{String=>Column}] target column => new value.
+      # @return [MergeIntoWriter]
+      def update(assignments) = @writer.add_action(:matched, :ACTION_TYPE_UPDATE, @condition, assignments)
+
+      # Delete the matched target row. @return [MergeIntoWriter]
+      def delete = @writer.add_action(:matched, :ACTION_TYPE_DELETE, @condition)
+    end
+
+    # Clause builder returned by {MergeIntoWriter#when_not_matched}.
+    class WhenNotMatched
+      def initialize(writer, condition)
+        @writer = writer
+        @condition = condition
+      end
+
+      # Insert the source row with all its columns. @return [MergeIntoWriter]
+      def insert_all = @writer.add_action(:not_matched, :ACTION_TYPE_INSERT_STAR, @condition)
+      alias insertAll insert_all
+
+      # @param assignments [Hash{String=>Column}] target column => value.
+      # @return [MergeIntoWriter]
+      def insert(assignments) = @writer.add_action(:not_matched, :ACTION_TYPE_INSERT, @condition, assignments)
+    end
+
+    # Clause builder returned by {MergeIntoWriter#when_not_matched_by_source}.
+    class WhenNotMatchedBySource
+      def initialize(writer, condition)
+        @writer = writer
+        @condition = condition
+      end
+
+      # @return [MergeIntoWriter]
+      def update_all = @writer.add_action(:not_matched_by_source, :ACTION_TYPE_UPDATE_STAR, @condition)
+      alias updateAll update_all
+
+      # @param assignments [Hash{String=>Column}] target column => new value.
+      # @return [MergeIntoWriter]
+      def update(assignments) = @writer.add_action(:not_matched_by_source, :ACTION_TYPE_UPDATE, @condition, assignments)
+
+      # Delete the unmatched target row. @return [MergeIntoWriter]
+      def delete = @writer.add_action(:not_matched_by_source, :ACTION_TYPE_DELETE, @condition)
+    end
+  end
+end

data/lib/spark_connect/na_functions.rb

@@ -12,6 +12,10 @@ module SparkConnect
   class DataFrameNaFunctions
     Proto = SparkConnect::Proto
 
+    # Marker for an omitted `value` argument of {#replace} (distinct from `nil`,
+    # which replaces with null).
+    NO_VALUE = Object.new.freeze
+
     # @param df [DataFrame]
     def initialize(df)
       @df = df
@@ -26,11 +30,17 @@ module SparkConnect
     # @return [DataFrame]
     def drop(how: :any, thresh: nil, subset: nil)
       cols = Array(subset).map(&:to_s)
-      min_non_nulls = thresh || (if how.to_sym == :all
-                                   1
-                                 else
-                                   (cols.empty? ? nil : cols.size)
-                                 end)
+      min_non_nulls =
+        case how&.to_sym
+        when :all then 1
+        when :any, nil then nil
+        else raise IllegalArgumentError, "how must be :any or :all, got #{how.inspect}"
+        end
+      unless thresh.nil?
+        raise IllegalArgumentError, "thresh must be an Integer, got #{thresh.class}" unless thresh.is_a?(Integer)
+
+        min_non_nulls = thresh # `thresh` overrides `how`, like PySpark
+      end
       nd = Proto::NADrop.new(input: @df.relation, cols: cols)
       nd.min_non_nulls = min_non_nulls if min_non_nulls
       @df.build(drop_na: nd)
@@ -46,10 +56,14 @@ module SparkConnect
     def fill(value, subset: nil)
       cols, values =
         if value.is_a?(Hash)
+          raise IllegalArgumentError, "fill value Hash must not be empty" if value.empty?
+
           [value.keys.map(&:to_s), value.values]
         else
-          [Array(subset).map(&:to_s), Array(subset).empty? ? [value] : Array(subset).map { value }]
+          # A single value applies to every column in `subset` (or all columns).
+          [Array(subset).map(&:to_s), [value]]
         end
+      values.each { |v| check_fill_value(v) }
       nf = Proto::NAFill.new(
         input: @df.relation, cols: cols, values: values.map { |v| na_literal(v) }
       )
@@ -64,16 +78,26 @@ module SparkConnect
     #   is not a Hash.
     # @param subset [Array<String>, nil]
     # @return [DataFrame]
-    def replace(to_replace, value = nil, subset: nil)
+    def replace(to_replace, value = NO_VALUE, subset: nil)
       mapping =
         if to_replace.is_a?(Hash)
           to_replace
         else
-          Array(to_replace).zip(Array(value)).to_h
+          raise IllegalArgumentError, "value is required when to_replace is not a Hash" if value.equal?(NO_VALUE)
+
+          olds = Array(to_replace)
+          news = value.is_a?(Array) ? value : Array.new(olds.size, value)
+          if olds.size != news.size
+            raise IllegalArgumentError, "to_replace and value must have the same length (#{olds.size} vs #{news.size})"
+          end
+
+          olds.zip(news).to_h
         end
+      # Like PySpark, numeric replacements are sent as doubles (Spark converts
+      # them per column type).
       replacements = mapping.map do |old, new_value|
         Proto::NAReplace::Replacement.new(
-          old_value: na_literal(old), new_value: na_literal(new_value)
+          old_value: replace_literal(old), new_value: replace_literal(new_value)
         )
       end
       nr = Proto::NAReplace.new(
@@ -94,5 +118,15 @@ module SparkConnect
       else Column.to_literal(value)
       end
     end
+
+    def replace_literal(value)
+      value.is_a?(Integer) ? Proto::Expression::Literal.new(double: value.to_f) : na_literal(value)
+    end
+
+    def check_fill_value(value)
+      return if [true, false].include?(value) || value.is_a?(Numeric) || value.is_a?(String)
+
+      raise IllegalArgumentError, "fill values must be Boolean, Numeric, or String, got #{value.class}"
+    end
   end
 end

data/lib/spark_connect/observation.rb

@@ -47,15 +47,36 @@ module SparkConnect
       result = @df.session.client.execute_plan(@df.relation)
       observed = result.observed_metrics.find { |m| m.name == @name } || result.observed_metrics.first
       return {} unless observed
+      # Spark 4.2+ reports a failure to compute the metrics (e.g. an error in a
+      # metric expression) in the metrics message instead of failing the query.
+      raise Observation.error_from(observed) if observed.has_root_error_idx?
 
       keys = observed.keys.to_a
-      decoded = observed.values.map { |lit| decode_literal(lit) }
-      keys.zip(decoded).to_h
+      keys = Array.new(observed.values.size) { |i| "observed_metric_#{i}" } if keys.empty?
+      keys.zip(observed.values.map { |lit| Column.literal_value(lit) }).to_h
     end
 
-    def decode_literal(literal)
-      kind = literal.literal_type
-      kind ? literal.public_send(kind) : nil
+    class << self
+      # Convert the error payload of an observed-metrics message into the
+      # matching {SparkConnectError} subclass.
+      # @api private
+      def error_from(observed)
+        errors = observed.errors.to_a
+        idx = observed.root_error_idx
+        return SparkConnectError.new("Observation error: invalid root_error_idx #{idx}") unless idx.between?(0, errors.size - 1)
+
+        root = errors[idx]
+        hierarchy = root.error_type_hierarchy.to_a
+        klass =
+          if hierarchy.any? { |c| c.end_with?("ParseException") } then ParseError
+          elsif hierarchy.any? { |c| c.end_with?("AnalysisException") } then AnalysisError
+          else SparkConnectError
+          end
+        throwable = root.spark_throwable
+        klass.new(root.message,
+                  error_class: throwable&.has_error_class? ? throwable.error_class : nil,
+                  sql_state: throwable&.has_sql_state? ? throwable.sql_state : nil)
+      end
     end
   end
 end

data/lib/spark_connect/pipelines.rb

@@ -70,6 +70,79 @@ module SparkConnect
                                             table_properties: table_properties, schema: schema)
     end
 
+    # Define a streaming table that flows (e.g. {#create_auto_cdc_flow} or
+    # {#define_flow}) can target, without a defining query. Mirrors PySpark's
+    # `create_streaming_table`.
+    # @return [String] the resolved output identifier.
+    def create_streaming_table(name, comment: nil, format: nil, partition_cols: [], clustering_columns: [],
+                               table_properties: {}, schema: nil)
+      define_table_output(name, :table, nil, comment: comment, format: format, partition_cols: partition_cols,
+                                             clustering_columns: clustering_columns, table_properties: table_properties,
+                                             schema: schema)
+    end
+
+    # Define an Auto CDC flow that applies the Change Data Capture events of
+    # `source` to the streaming table `target` (Spark 4.2+). Mirrors PySpark's
+    # `create_auto_cdc_flow`.
+    #
+    # @example
+    #   pipe.create_streaming_table("customers")
+    #   pipe.create_auto_cdc_flow(target: "customers", source: "customers_cdc", keys: ["id"],
+    #                             sequence_by: "ts", apply_as_deletes: "op = 'DELETE'", stored_as_scd_type: 2)
+    #
+    # @param target [String] the target streaming table.
+    # @param source [String] the CDC source to stream from.
+    # @param keys [Array<String, Column>] the columns identifying a row (unqualified).
+    # @param sequence_by [String, Column] the ordering expression (a SQL String or Column).
+    # @param apply_as_deletes [String, Column, nil] the condition marking a delete event.
+    # @param column_list [Array<String, Column>, nil] columns to keep in the target.
+    # @param except_column_list [Array<String, Column>, nil] columns to drop from the target.
+    # @param stored_as_scd_type [Integer, String, nil] 1 or 2 (server default when nil).
+    # @param name [String, nil] the flow name (defaults to `target`).
+    # @param track_history_column_list [Array<String, Column>, nil] SCD2 only: columns whose change opens a new history record.
+    # @param track_history_except_column_list [Array<String, Column>, nil] SCD2 only: columns excluded from history tracking.
+    # @param sql_conf [Hash{String=>String}] SQL configs applied to this flow.
+    # @return [String] the resolved flow name.
+    def create_auto_cdc_flow(target:, source:, keys:, sequence_by:, apply_as_deletes: nil, column_list: nil,
+                             except_column_list: nil, stored_as_scd_type: nil, name: nil,
+                             track_history_column_list: nil, track_history_except_column_list: nil, sql_conf: {})
+      %i[target source].zip([target, source]).each do |arg, value|
+        raise IllegalArgumentError, "#{arg} must be a String, got #{value.class}" unless value.is_a?(String)
+      end
+      raise IllegalArgumentError, "name must be a String, got #{name.class}" unless name.nil? || name.is_a?(String)
+
+      key_cols = cdc_columns("keys", keys)
+      columns = cdc_columns("column_list", column_list)
+      except = cdc_columns("except_column_list", except_column_list)
+      track = cdc_columns("track_history_column_list", track_history_column_list)
+      track_except = cdc_columns("track_history_except_column_list", track_history_except_column_list)
+      reject_together("column_list", columns, "except_column_list", except)
+      reject_together("track_history_column_list", track, "track_history_except_column_list", track_except)
+      unless stored_as_scd_type.nil? || %w[1 2].include?(stored_as_scd_type.to_s)
+        raise IllegalArgumentError, "stored_as_scd_type must be 1 or 2, got #{stored_as_scd_type.inspect}"
+      end
+      if stored_as_scd_type.to_s != "2" && (track&.any? || track_except&.any?)
+        raise IllegalArgumentError,
+              "track_history_column_list/track_history_except_column_list require stored_as_scd_type 2"
+      end
+
+      details = PC::DefineFlow::AutoCdcFlowDetails.new(
+        source: source, keys: key_cols.map(&:to_expr), sequence_by: cdc_expression("sequence_by", sequence_by).to_expr,
+        column_list: (columns || []).map(&:to_expr), except_column_list: (except || []).map(&:to_expr),
+        track_history_column_list: (track || []).map(&:to_expr),
+        track_history_except_column_list: (track_except || []).map(&:to_expr)
+      )
+      details.apply_as_deletes = cdc_expression("apply_as_deletes", apply_as_deletes).to_expr unless apply_as_deletes.nil?
+      details.stored_as_scd_type = :"SCD_TYPE_#{stored_as_scd_type}" unless stored_as_scd_type.nil?
+      flow_name = (name || target).to_s
+      flow = PC::DefineFlow.new(
+        dataflow_graph_id: @graph_id, flow_name: flow_name, target_dataset_name: target,
+        sql_conf: stringify(sql_conf), auto_cdc_flow_details: details
+      )
+      result = dispatch(PC.new(define_flow: flow))
+      identifier_string(result.pipeline_command_result&.define_flow_result&.resolved_identifier) || flow_name
+    end
+
     # Define a materialized view and the flow that populates it.
     # @return [String]
     def create_materialized_view(name, df = nil, comment: nil, format: nil, partition_cols: [],
@@ -195,6 +268,34 @@ module SparkConnect
       identifier_string(result.pipeline_command_result&.define_output_result&.resolved_identifier) || name.to_s
     end
 
+    def cdc_columns(arg, cols)
+      return nil if cols.nil?
+      raise IllegalArgumentError, "#{arg} must be an Array of Strings or Columns, got #{cols.class}" unless cols.is_a?(Array)
+
+      cols.map do |c|
+        case c
+        when Column then c
+        when String then Functions.col(c)
+        else raise IllegalArgumentError, "#{arg} must contain only Strings or Columns, got #{c.class}"
+        end
+      end
+    end
+
+    def cdc_expression(arg, value)
+      case value
+      when Column then value
+      when String then Functions.expr(value)
+      else raise IllegalArgumentError, "#{arg} must be a SQL String or Column, got #{value.class}"
+      end
+    end
+
+    # An empty list is the same as an omitted one, so only non-empty lists clash.
+    def reject_together(include_arg, include_cols, except_arg, except_cols)
+      return unless include_cols&.any? && except_cols&.any?
+
+      raise IllegalArgumentError, "#{include_arg} and #{except_arg} cannot be set together"
+    end
+
     def apply_schema(details, schema)
       if schema.is_a?(Types::DataType)
         details.schema_data_type = schema.to_proto

data/lib/spark_connect/plan.rb

@@ -35,5 +35,57 @@ module SparkConnect
     def command_plan(command)
       Proto::Plan.new(command: command)
     end
+
+    # When the expressions of `relation` (not of its child relations, which
+    # were wrapped when they were built) reference subquery DataFrames
+    # ({DataFrame#scalar}, {DataFrame#exists}, {Column#isin} with a DataFrame,
+    # {DataFrame#as_table}), wrap it in a `WithRelations` that carries those
+    # subquery plans, like PySpark:
+    #
+    #     with_relations [new plan id]
+    #       root: relation
+    #       references: [subquery plans...]
+    #
+    # @param session [SparkSession] the session that registered the subqueries.
+    # @param relation [Spark::Connect::Relation]
+    # @return [Spark::Connect::Relation]
+    def with_subquery_references(session, relation)
+      ids = []
+      body = relation.rel_type && relation.public_send(relation.rel_type)
+      collect_subquery_plan_ids(body, ids) if body
+      return relation if ids.empty?
+
+      references = ids.uniq.map do |id|
+        session.subquery_relation(id) or
+          raise IllegalArgumentError, "Subquery DataFrame (plan id #{id}) does not belong to this SparkSession"
+      end
+      Proto::Relation.new(
+        common: Proto::RelationCommon.new(plan_id: session.next_plan_id),
+        with_relations: Proto::WithRelations.new(root: relation, references: references)
+      )
+    end
+
+    # Collect into `acc` the plan ids of the subquery expressions in `msg`,
+    # without descending into child relations.
+    # @api private
+    def collect_subquery_plan_ids(msg, acc)
+      return acc if msg.is_a?(Proto::Relation)
+
+      acc << msg.plan_id if msg.is_a?(Proto::SubqueryExpression)
+      msg.class.descriptor.each do |field|
+        next unless field.type == :message
+
+        value = msg[field.name]
+        case value
+        when Google::Protobuf::RepeatedField then value.each { |v| collect_subquery_plan_ids(v, acc) }
+        when Google::Protobuf::Map
+          # Google::Protobuf::Map has no #each_value; iterate its message values.
+          value.values.grep(Google::Protobuf::MessageExts).each { |v| collect_subquery_plan_ids(v, acc) }
+        when nil then nil
+        else collect_subquery_plan_ids(value, acc)
+        end
+      end
+      acc
+    end
   end
 end