spark-connect 0.2.0

data/lib/spark_connect/session.rb

@@ -0,0 +1,317 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+module SparkConnect
+  # The entry point to programming Spark with the DataFrame API over Spark
+  # Connect. Create one with the {Builder}:
+  #
+  # @example
+  #   spark = SparkConnect::SparkSession.builder
+  #                                     .remote("sc://localhost:15002")
+  #                                     .app_name("my-app")
+  #                                     .get_or_create
+  #
+  # A session owns the underlying {SparkConnectClient}, a monotonic plan-id
+  # allocator (so each relation is uniquely identifiable to the server), and the
+  # {RuntimeConfig} and {Catalog} facades.
+  class SparkSession
+    Proto = SparkConnect::Proto
+
+    # @return [SparkConnectClient]
+    attr_reader :client
+
+    # @param client [SparkConnectClient]
+    def initialize(client)
+      @client = client
+      @plan_id = -1
+      @conf = RuntimeConfig.new(client)
+    end
+
+    class << self
+      # @return [Builder] a new session builder.
+      def builder
+        Builder.new
+      end
+
+      # The currently active session set by {#set_active} / {Builder#get_or_create}.
+      # @return [SparkSession, nil]
+      attr_accessor :active
+
+      # @api private
+    end
+
+    # Allocate the next unique plan id. Used by {PlanBuilder.relation}.
+    # @return [Integer]
+    def next_plan_id
+      @plan_id += 1
+    end
+
+    # @return [String] the client session id (UUID).
+    def session_id
+      @client.session_id
+    end
+
+    # Create a {DataFrame} with a single `id` column over the given integer range.
+    #
+    # @overload range(end_)
+    # @overload range(start, end_, step = 1, num_partitions = nil)
+    # @return [DataFrame]
+    def range(start, end_ = nil, step = 1, num_partitions = nil)
+      if end_.nil?
+        end_ = start
+        start = 0
+      end
+      r = Proto::Range.new(start: start, end: end_, step: step)
+      r.num_partitions = num_partitions if num_partitions
+      DataFrame.new(self, PlanBuilder.relation(self, range: r))
+    end
+
+    # Execute a SQL query and return a lazy {DataFrame} over its result.
+    #
+    # @param query [String]
+    # @param args [Hash{String=>Object}, Array<Object>, nil] named or positional
+    #   parameters bound into the query.
+    # @return [DataFrame]
+    def sql(query, args = nil)
+      sql = Proto::SQL.new(query: query)
+      case args
+      when Hash
+        args.each { |k, v| sql.named_arguments[k.to_s] = Column.to_col(v).to_expr }
+      when Array
+        sql.pos_arguments += args.map { |v| Column.to_col(v).to_expr }
+      end
+      DataFrame.new(self, PlanBuilder.relation(self, sql: sql))
+    end
+
+    # Return a {DataFrame} reading the named table or view.
+    #
+    # @param name [String]
+    # @return [DataFrame]
+    def table(name)
+      read.table(name)
+    end
+
+    # @return [DataFrameReader] interface for loading external data.
+    def read
+      DataFrameReader.new(self)
+    end
+
+    # @return [DataStreamReader] interface for loading a streaming DataFrame.
+    def read_stream
+      DataStreamReader.new(self)
+    end
+    alias readStream read_stream
+
+    # @return [StreamingQueryManager] the manager for this session's streaming queries.
+    def streams
+      StreamingQueryManager.new(self)
+    end
+
+    # Create a new Spark Declarative Pipeline (a dataflow graph) in this session.
+    #
+    # @param default_catalog [String, nil]
+    # @param default_database [String, nil]
+    # @param sql_conf [Hash{String=>String}] SQL configs applied to all flows.
+    # @return [Pipeline]
+    def pipeline(default_catalog: nil, default_database: nil, sql_conf: {})
+      Pipeline.new(self, default_catalog: default_catalog, default_database: default_database, sql_conf: sql_conf)
+    end
+
+    # Build a {DataFrame} from local Ruby data.
+    #
+    # @param data [Array<Hash>, Array<Array>, Array<Row>]
+    # @param schema [Types::StructType, Array<String>, String, nil] an explicit
+    #   schema, a list of column names, a DDL string, or `nil` to infer.
+    # @return [DataFrame]
+    def create_data_frame(data, schema = nil)
+      data = data.to_a
+      struct = resolve_schema(data, schema)
+      arrow_bytes = ArrowConverter.from_rows(data, struct)
+      local = Proto::LocalRelation.new(data: arrow_bytes, schema: struct.simple_string.sub(/\Astruct</, "").sub(/>\z/, ""))
+      DataFrame.new(self, PlanBuilder.relation(self, local_relation: local))
+    end
+    alias create_dataframe create_data_frame
+    alias createDataFrame create_data_frame
+
+    # @return [RuntimeConfig] runtime configuration facade.
+    attr_reader :conf
+
+    # @return [Catalog] the catalog facade (databases, tables, functions, cache).
+    def catalog
+      @catalog ||= Catalog.new(self)
+    end
+
+    # @return [String] the Spark version reported by the server.
+    def version
+      @client.analyze(spark_version: Proto::AnalyzePlanRequest::SparkVersion.new).spark_version.version
+    end
+
+    # Make this the active/default session.
+    # @return [self]
+    def set_active
+      SparkSession.active = self
+      self
+    end
+
+    # Start a brand-new session against the same endpoint (independent
+    # server-side session id, configuration, and temporary views).
+    # @return [SparkSession]
+    def new_session
+      SparkSession.new(SparkConnectClient.new(@client.channel_builder))
+    end
+
+    # Interrupt all operations running in this session.
+    # @return [Array<String>] the ids of the interrupted operations.
+    def interrupt_all
+      @client.interrupt(type: :all).interrupted_ids.to_a
+    end
+
+    # Interrupt all operations tagged with `tag` (see {#add_tag}).
+    # @return [Array<String>]
+    def interrupt_tag(tag)
+      @client.interrupt(type: :tag, value: tag.to_s).interrupted_ids.to_a
+    end
+
+    # Interrupt a single operation by id.
+    # @return [Array<String>]
+    def interrupt_operation(operation_id)
+      @client.interrupt(type: :operation_id, value: operation_id.to_s).interrupted_ids.to_a
+    end
+
+    # Add an operation tag applied to all subsequent executions in this session.
+    # @return [void]
+    def add_tag(tag)
+      @client.add_tag(tag)
+    end
+
+    # Remove a previously added operation tag. @return [void]
+    def remove_tag(tag)
+      @client.remove_tag(tag)
+    end
+
+    # @return [Array<String>] the currently active operation tags.
+    def get_tags
+      @client.tags.dup
+    end
+
+    # Remove all operation tags. @return [void]
+    def clear_tags
+      @client.clear_tags
+    end
+
+    # Release the server-side session and stop the client.
+    # @return [void]
+    def stop
+      @client.release_session
+      SparkSession.active = nil if SparkSession.active.equal?(self)
+      nil
+    end
+
+    # @api private
+    def create_data_frame_from_relation(relation)
+      DataFrame.new(self, relation)
+    end
+
+    private
+
+    def resolve_schema(data, schema)
+      case schema
+      when Types::StructType then schema
+      when String then parse_ddl_schema(schema)
+      when Array then infer_schema(data, names: schema.map(&:to_s))
+      when nil then infer_schema(data)
+      else
+        raise IllegalArgumentError, "Unsupported schema: #{schema.inspect}"
+      end
+    end
+
+    def parse_ddl_schema(ddl)
+      # Ask the server to parse the DDL into a concrete schema.
+      proto = @client.analyze(ddl_parse: Proto::AnalyzePlanRequest::DDLParse.new(ddl_string: ddl)).ddl_parse.parsed
+      Types.from_proto(proto)
+    end
+
+    def infer_schema(data, names: nil)
+      raise IllegalArgumentError, "Cannot infer schema from empty data; pass a schema" if data.empty?
+
+      first = data.first
+      case first
+      when Hash
+        keys = first.keys.map(&:to_s)
+        Types::StructType.new(keys.map.with_index do |k, i|
+          Types::StructField.new(names ? names[i] : k, column_type(data, k, i), nullable: true)
+        end)
+      when Row
+        Types::StructType.new(first.fields.map.with_index do |k, i|
+          Types::StructField.new(names ? names[i] : k, column_type(data, k, i), nullable: true)
+        end)
+      when Array
+        Types::StructType.new(first.each_index.map do |i|
+          Types::StructField.new(names ? names[i] : "_#{i + 1}", column_type(data, nil, i), nullable: true)
+        end)
+      else
+        raise IllegalArgumentError, "Cannot infer schema from rows of type #{first.class}"
+      end
+    end
+
+    def column_type(data, key, index)
+      sample = data.map { |row| ArrowConverter.extract_value(row, key, index) }.find { |v| !v.nil? }
+      Column.infer_type(sample)
+    end
+  end
+
+  # Fluent builder for {SparkSession}. Returned by {SparkSession.builder}.
+  class SparkSession
+    class Builder
+      def initialize
+        @options = {}
+        @remote = nil
+      end
+
+      # Set the connection string (`sc://...`).
+      # @return [self]
+      def remote(url)
+        @remote = url
+        self
+      end
+
+      # Set the application name.
+      # @return [self]
+      def app_name(name)
+        @options["spark.app.name"] = name
+        self
+      end
+
+      # Set an arbitrary configuration option to apply after connecting.
+      # @return [self]
+      def config(key, value)
+        @options[key.to_s] = value
+        self
+      end
+
+      # Build (or reuse the active) {SparkSession}.
+      # @return [SparkSession]
+      def get_or_create
+        existing = SparkSession.active
+        return existing if existing
+
+        session = create
+        SparkSession.active = session
+        session
+      end
+      alias getOrCreate get_or_create
+
+      # Always build a new {SparkSession}.
+      # @return [SparkSession]
+      def create
+        url = @remote || ENV["SPARK_REMOTE"] || "sc://localhost:15002"
+        client = SparkConnectClient.new(ChannelBuilder.new(url))
+        session = SparkSession.new(client)
+        @options.each { |k, v| session.conf.set(k, v) unless k == "spark.app.name" }
+        session
+      end
+      alias build create
+    end
+  end
+end

data/lib/spark_connect/stat_functions.rb

@@ -0,0 +1,109 @@
+# frozen_string_literal: true
+
+module SparkConnect
+  # Statistical helpers, returned by {DataFrame#stat}. Mirrors PySpark's
+  # `DataFrame.stat` (`DataFrameStatFunctions`).
+  #
+  # @example
+  #   df.stat.corr("x", "y")
+  #   df.stat.approx_quantile("x", [0.25, 0.5, 0.75], 0.01)
+  #   df.stat.crosstab("a", "b").show
+  class DataFrameStatFunctions
+    Proto = SparkConnect::Proto
+
+    # @param df [DataFrame]
+    def initialize(df)
+      @df = df
+    end
+
+    # Sample covariance of two columns.
+    # @return [Float]
+    def cov(col1, col2)
+      scalar(@df.build(cov: Proto::StatCov.new(input: @df.relation, col1: col1.to_s, col2: col2.to_s)))
+    end
+
+    # Correlation of two columns (`method` is `"pearson"`).
+    # @return [Float]
+    def corr(col1, col2, method = "pearson")
+      rel = Proto::StatCorr.new(input: @df.relation, col1: col1.to_s, col2: col2.to_s, method: method)
+      scalar(@df.build(corr: rel))
+    end
+
+    # Contingency table (cross-tabulation) of two columns.
+    # @return [DataFrame]
+    def crosstab(col1, col2)
+      @df.build(crosstab: Proto::StatCrosstab.new(input: @df.relation, col1: col1.to_s, col2: col2.to_s))
+    end
+
+    # Frequent items in the given columns.
+    #
+    # @param cols [Array<String>]
+    # @param support [Float]
+    # @return [DataFrame]
+    def freq_items(cols, support = 0.01)
+      rel = Proto::StatFreqItems.new(input: @df.relation, cols: Array(cols).map(&:to_s), support: support)
+      @df.build(freq_items: rel)
+    end
+
+    # Approximate quantiles of numeric columns.
+    #
+    # @param cols [String, Array<String>]
+    # @param probabilities [Array<Float>] values in 0.0..1.0.
+    # @param relative_error [Float]
+    # @return [Array<Float>, Array<Array<Float>>] one list per column.
+    def approx_quantile(cols, probabilities, relative_error)
+      single = !cols.is_a?(Array)
+      rel = Proto::StatApproxQuantile.new(
+        input: @df.relation, cols: Array(cols).map(&:to_s),
+        probabilities: probabilities, relative_error: relative_error
+      )
+      row = @df.build(approx_quantile: rel).collect.first
+      result = row.to_a
+      single ? result.first : result
+    end
+
+    # Stratified sample without replacement, keyed by `col`.
+    #
+    # @param col [String, Column]
+    # @param fractions [Hash{Object=>Float}] per-stratum sampling fraction.
+    # @param seed [Integer, nil]
+    # @return [DataFrame]
+    def sample_by(col, fractions, seed = nil)
+      col_expr = (col.is_a?(Column) ? col : Functions.col(col.to_s)).to_expr
+      frac = fractions.map do |stratum, fraction|
+        Proto::StatSampleBy::Fraction.new(stratum: Column.to_literal(stratum), fraction: fraction)
+      end
+      rel = Proto::StatSampleBy.new(input: @df.relation, col: col_expr, fractions: frac)
+      rel.seed = seed if seed
+      @df.build(sample_by: rel)
+    end
+
+    private
+
+    def scalar(df)
+      row = df.collect.first
+      row&.[](0)
+    end
+  end
+
+  # Reopen {DataFrame} to add the `describe`/`summary` actions, which are
+  # naturally statistical and share the Stat* relations.
+  class DataFrame
+    # Basic descriptive statistics (count, mean, stddev, min, max) per column.
+    #
+    # @param cols [Array<String>] columns to describe (all when empty).
+    # @return [DataFrame]
+    def describe(*cols)
+      build(describe: Proto::StatDescribe.new(input: @relation, cols: cols.flatten.map(&:to_s)))
+    end
+
+    # Configurable summary statistics.
+    #
+    # @param statistics [Array<String>] e.g. `"count"`, `"mean"`, `"stddev"`,
+    #   `"min"`, `"25%"`, `"50%"`, `"75%"`, `"max"`.
+    # @return [DataFrame]
+    def summary(*statistics)
+      build(summary: Proto::StatSummary.new(input: @relation, statistics: statistics.flatten.map(&:to_s)))
+    end
+  end
+end