spark-connect 0.3.0 → 0.4.0

data/lib/spark_connect/storage_level.rb

@@ -0,0 +1,79 @@
+# frozen_string_literal: true
+
+module SparkConnect
+  # How a cached {DataFrame} (or table) is stored: in memory and/or on disk,
+  # off-heap, deserialized, and with how many replicas. Mirrors PySpark's
+  # `StorageLevel`; use the predefined constants (`StorageLevel::MEMORY_AND_DISK`, ...).
+  #
+  # @example
+  #   df.persist(SparkConnect::StorageLevel::DISK_ONLY)
+  #   df.storage_level  #=> StorageLevel(true, false, false, false, 1)
+  class StorageLevel
+    Proto = SparkConnect::Proto
+
+    # @return [Boolean]
+    attr_reader :use_disk, :use_memory, :use_off_heap, :deserialized
+    # @return [Integer]
+    attr_reader :replication
+
+    def initialize(use_disk, use_memory, use_off_heap, deserialized, replication = 1)
+      @use_disk = use_disk
+      @use_memory = use_memory
+      @use_off_heap = use_off_heap
+      @deserialized = deserialized
+      @replication = replication
+    end
+
+    # @return [Spark::Connect::StorageLevel]
+    def to_proto
+      Proto::StorageLevel.new(use_disk: use_disk, use_memory: use_memory, use_off_heap: use_off_heap,
+                              deserialized: deserialized, replication: replication)
+    end
+
+    # @param proto [Spark::Connect::StorageLevel]
+    # @return [StorageLevel]
+    def self.from_proto(proto)
+      new(proto.use_disk, proto.use_memory, proto.use_off_heap, proto.deserialized, proto.replication)
+    end
+
+    # @return [Boolean] whether `other` is a StorageLevel with the same flags.
+    def ==(other)
+      other.is_a?(StorageLevel) && to_a == other.to_a
+    end
+    alias eql? ==
+
+    # @return [Integer]
+    def hash = to_a.hash
+
+    # @return [Array] `[use_disk, use_memory, use_off_heap, deserialized, replication]`.
+    def to_a = [use_disk, use_memory, use_off_heap, deserialized, replication]
+
+    # @return [String] e.g. `"StorageLevel(true, true, false, false, 1)"`.
+    def to_s
+      "StorageLevel(#{to_a.join(', ')})"
+    end
+    alias inspect to_s
+
+    # Not stored.
+    NONE = new(false, false, false, false)
+    # On disk only.
+    DISK_ONLY = new(true, false, false, false)
+    # On disk only, replicated on two nodes.
+    DISK_ONLY_2 = new(true, false, false, false, 2)
+    # On disk only, replicated on three nodes.
+    DISK_ONLY_3 = new(true, false, false, false, 3)
+    # In memory only (serialized).
+    MEMORY_ONLY = new(false, true, false, false)
+    # In memory only, replicated on two nodes.
+    MEMORY_ONLY_2 = new(false, true, false, false, 2)
+    # In memory (serialized), spilling to disk.
+    MEMORY_AND_DISK = new(true, true, false, false)
+    # In memory and disk, replicated on two nodes.
+    MEMORY_AND_DISK_2 = new(true, true, false, false, 2)
+    # Off-heap memory.
+    OFF_HEAP = new(true, true, true, false, 1)
+    # In memory (deserialized), spilling to disk; the default of {DataFrame#cache}
+    # and {DataFrame#persist}.
+    MEMORY_AND_DISK_DESER = new(true, true, false, true)
+  end
+end

data/lib/spark_connect/streaming.rb

@@ -11,12 +11,15 @@ module SparkConnect
   class DataStreamReader
     Proto = SparkConnect::Proto
 
+    SOURCE_NAME_PATTERN = /\A[a-zA-Z0-9_]+\z/
+
     # @param session [SparkSession]
     def initialize(session)
       @session = session
       @format = nil
       @schema = nil
       @options = {}
+      @source_name = nil
     end
 
     # @return [self] set the streaming source format (`"rate"`, `"kafka"`, ...).
@@ -27,19 +30,37 @@ module SparkConnect
 
     # @return [self] set the input schema (a {Types::StructType} or DDL string).
     def schema(schema)
-      @schema = schema.is_a?(Types::StructType) ? schema.simple_string : schema.to_s
+      # Like PySpark, a StructType is sent as its JSON form.
+      @schema = schema.is_a?(Types::StructType) ? schema.json : schema.to_s
       self
     end
 
-    # @return [self] set a single source option.
+    # @return [self] set a single source option (a `nil` value is ignored).
     def option(key, value)
-      @options[key.to_s] = value.to_s
+      @options[key.to_s] = value.to_s unless value.nil?
       self
     end
 
     # @return [self] set multiple source options.
     def options(opts)
-      opts.each { |k, v| @options[k.to_s] = v.to_s }
+      opts.each { |k, v| option(k, v) }
+      self
+    end
+
+    # Name this streaming source (Spark 4.2+). The name identifies the source
+    # in checkpoint metadata, so sources can be added, removed, or reordered
+    # without breaking checkpoint compatibility.
+    #
+    # @param source_name [String] letters, digits, and underscores only.
+    # @return [self]
+    def name(source_name)
+      raise IllegalArgumentError, "source name must be a String, got #{source_name.class}" unless source_name.is_a?(String)
+      unless SOURCE_NAME_PATTERN.match?(source_name)
+        raise IllegalArgumentError,
+              "Invalid streaming source name #{source_name.inspect}: only ASCII letters, digits, and underscores are allowed"
+      end
+
+      @source_name = source_name
       self
     end
 
@@ -51,9 +72,23 @@ module SparkConnect
       ds = Proto::Read::DataSource.new(options: @options, paths: path ? [path.to_s] : [])
       ds.format = @format if @format
       ds.schema = @schema if @schema
+      ds.source_name = @source_name if @source_name
       stream_relation(data_source: ds)
     end
 
+    # The row-level changes (Change Data Capture) of a table as a streaming
+    # DataFrame (Spark 4.2+). Configure the start with options such as
+    # `startingVersion` / `startingTimestamp`.
+    #
+    # @param table_name [String]
+    # @return [DataFrame]
+    def changes(table_name)
+      raise AnalysisError.new("User specified schema not supported with `changes`.", error_class: "_LEGACY_ERROR_TEMP_1189") if @schema
+
+      rc = Proto::RelationChanges.new(unparsed_identifier: table_name.to_s, options: @options, is_streaming: true)
+      DataFrame.new(@session, PlanBuilder.relation(@session, relation_changes: rc))
+    end
+
     # Load a streaming DataFrame from a registered table.
     # @return [DataFrame]
     def table(name)
@@ -66,6 +101,8 @@ module SparkConnect
     def parquet(path) = format("parquet").load(path)
     def orc(path) = format("orc").load(path)
     def text(path) = format("text").load(path)
+    # @return [DataFrame] a streaming XML source (set `option("rowTag", ...)`).
+    def xml(path) = format("xml").load(path)
 
     private
 
@@ -78,9 +115,8 @@ module SparkConnect
   # Writes a streaming {DataFrame} to a streaming sink and starts the query.
   # Returned by {DataFrame#write_stream}. Mirrors PySpark's `DataStreamWriter`.
   #
-  # `foreach`/`foreach_batch` are intentionally unsupported: they require
-  # user-defined functions, whose Spark Connect protobuf definitions are not yet
-  # finalized.
+  # `foreach`/`foreach_batch` are intentionally unsupported: they run a
+  # user-defined function on the server, which this client cannot ship.
   #
   # @example
   #   query = df.write_stream
@@ -104,6 +140,7 @@ module SparkConnect
       @trigger = nil
       @path = nil
       @table = nil
+      @clustering = []
     end
 
     # @return [self] set the sink format (`"console"`, `"memory"`, `"kafka"`, ...).
@@ -118,15 +155,15 @@ module SparkConnect
       self
     end
 
-    # @return [self] set a single sink option.
+    # @return [self] set a single sink option (a `nil` value is ignored).
     def option(key, value)
-      @options[key.to_s] = value.to_s
+      @options[key.to_s] = value.to_s unless value.nil?
       self
     end
 
     # @return [self] set multiple sink options.
     def options(opts)
-      opts.each { |k, v| @options[k.to_s] = v.to_s }
+      opts.each { |k, v| option(k, v) }
       self
     end
 
@@ -135,6 +172,27 @@ module SparkConnect
       @partitioning = cols.flatten.map(&:to_s)
       self
     end
+    alias partitionBy partition_by
+
+    # @return [self] cluster the output table by these columns.
+    def cluster_by(*cols)
+      @clustering = cols.flatten.map(&:to_s)
+      self
+    end
+    alias clusterBy cluster_by
+
+    # Not supported: `foreach` runs a user-defined function on the server.
+    # @raise [NotImplementedError]
+    def foreach(*_args)
+      raise NotImplementedError, "foreach sinks require user-defined functions, which the Ruby client does not support"
+    end
+
+    # Not supported: `foreach_batch` runs a user-defined function on the server.
+    # @raise [NotImplementedError]
+    def foreach_batch(*_args)
+      raise NotImplementedError, "foreachBatch sinks require user-defined functions, which the Ruby client does not support"
+    end
+    alias foreachBatch foreach_batch
 
     # @return [self] name the streaming query (required by the memory sink).
     def query_name(name)
@@ -148,22 +206,35 @@ module SparkConnect
     # @param once [Boolean, nil] process all available data once and stop.
     # @param available_now [Boolean, nil] process all available data in (possibly) multiple batches, then stop.
     # @param continuous [String, nil] continuous-processing checkpoint interval.
+    # @param real_time [String, nil] real-time mode with the given batch
+    #   duration, e.g. `"5 minutes"` (Spark 4.2+; supported sources/sinks only).
     # @return [self]
-    def trigger(processing_time: nil, once: nil, available_now: nil, continuous: nil)
+    def trigger(processing_time: nil, once: nil, available_now: nil, continuous: nil, real_time: nil)
+      given = { processing_time: processing_time, once: once, available_now: available_now,
+                continuous: continuous, real_time: real_time, }.compact
+      raise IllegalArgumentError, "trigger requires exactly one of #{given.keys.inspect}" if given.size > 1
+      raise IllegalArgumentError, "trigger requires one of processing_time, once, available_now, continuous, real_time" if given.empty?
+
+      key, value = given.first
       @trigger =
-        if processing_time then [:processing_time_interval, processing_time.to_s]
-        elsif once then [:once, true]
-        elsif available_now then [:available_now, true]
-        elsif continuous then [:continuous_checkpoint_interval, continuous.to_s]
+        case key
+        when :processing_time then [:processing_time_interval, non_empty!(key, value)]
+        when :once then [:once, true!(key, value)]
+        when :available_now then [:available_now, true!(key, value)]
+        when :continuous then [:continuous_checkpoint_interval, non_empty!(key, value)]
+        when :real_time then [:real_time_batch_duration, non_empty!(key, value)]
         end
       self
     end
 
-    # Start the streaming query to a file/data path.
+    # Start the streaming query to a file/data path. The keyword arguments are
+    # shortcuts for {#format}, {#output_mode}, {#partition_by}, {#query_name},
+    # and {#options}.
     #
     # @param path [String, nil]
     # @return [StreamingQuery]
-    def start(path = nil)
+    def start(path = nil, format: nil, output_mode: nil, partition_by: nil, query_name: nil, **options)
+      apply_start_options(format, output_mode, partition_by, query_name, options)
       @path = path if path
       run
     end
@@ -172,7 +243,8 @@ module SparkConnect
     #
     # @param name [String]
     # @return [StreamingQuery]
-    def to_table(name)
+    def to_table(name, format: nil, output_mode: nil, partition_by: nil, query_name: nil, **options)
+      apply_start_options(format, output_mode, partition_by, query_name, options)
       @table = name.to_s
       run
     end
@@ -180,10 +252,30 @@ module SparkConnect
 
     private
 
+    def apply_start_options(format, output_mode, partition_by, query_name, options)
+      self.format(format) if format
+      self.output_mode(output_mode) if output_mode
+      self.partition_by(partition_by) if partition_by
+      self.query_name(query_name) if query_name
+      self.options(options) unless options.empty?
+    end
+
+    def non_empty!(key, value)
+      raise IllegalArgumentError, "trigger #{key} must be a non-empty String" unless value.is_a?(String) && !value.strip.empty?
+
+      value.strip
+    end
+
+    def true!(key, value)
+      raise IllegalArgumentError, "trigger #{key} must be true" unless value == true
+
+      true
+    end
+
     def run
       op = WSO.new(
         input: @df.relation, format: @format || "", options: @options,
-        partitioning_column_names: @partitioning
+        partitioning_column_names: @partitioning, clustering_column_names: @clustering
       )
       op.output_mode = @output_mode if @output_mode
       op.query_name = @query_name if @query_name
@@ -195,6 +287,9 @@ module SparkConnect
       wsr = result.write_stream_result
       raise SparkConnectError, "Server did not return a streaming query handle" unless wsr
 
+      if wsr.query_started_event_json && !wsr.query_started_event_json.empty?
+        @df.session.streams.post_event(StreamingQueryListener.event_from_json(:started, wsr.query_started_event_json))
+      end
       StreamingQuery.new(@df.session, wsr.query_id, wsr.name)
     end
   end
@@ -209,6 +304,7 @@ module SparkConnect
     attr_reader :id
     # @return [String] the run id (unique per start).
     attr_reader :run_id
+    alias runId run_id
     # @return [String, nil] the query name, if one was set.
     attr_reader :name
 
@@ -239,16 +335,20 @@ module SparkConnect
     def active?
       status["isActive"]
     end
+    alias is_active active?
+    alias isActive active?
 
     # @return [Array<Hash>] parsed JSON progress objects for recent micro-batches.
     def recent_progress
       command(recent_progress: true).recent_progress.recent_progress_json.map { |j| JSON.parse(j) }
     end
+    alias recentProgress recent_progress
 
     # @return [Hash, nil] the most recent progress object, if any.
     def last_progress
       command(last_progress: true).recent_progress.recent_progress_json.map { |j| JSON.parse(j) }.last
     end
+    alias lastProgress last_progress
 
     # Block until the query terminates, or until `timeout_ms` elapses.
     #
@@ -266,6 +366,7 @@ module SparkConnect
       command(process_all_available: true)
       nil
     end
+    alias processAllAvailable process_all_available
 
     # Stop the query. @return [void]
     def stop
@@ -276,7 +377,10 @@ module SparkConnect
     # @return [String, nil] the query's exception message, if it has failed.
     def exception
       result = command(exception: true).exception
-      result.exception_message && result.exception_message.empty? ? nil : result.exception_message
+      # The server omits the result (or its message) when the query has not failed.
+      return nil if result.nil? || !result.has_exception_message? || result.exception_message.empty?
+
+      result.exception_message
     end
 
     # @return [String] the query's execution plan.
@@ -306,6 +410,44 @@ module SparkConnect
     # @param session [SparkSession]
     def initialize(session)
       @session = session
+      @listener_bus = StreamingQueryListenerBus.new(session)
+    end
+
+    # Register a {StreamingQueryListener} (or any object responding to some of
+    # `on_query_started`, `on_query_progress`, `on_query_idle`,
+    # `on_query_terminated`, or their camelCase PySpark names). Events are
+    # delivered on a background thread, like PySpark's client-side listener bus.
+    #
+    # @param listener [Object]
+    # @return [void]
+    def add_listener(listener)
+      @listener_bus.add(listener)
+      nil
+    end
+    alias addListener add_listener
+
+    # Unregister a listener added with {#add_listener}. @return [void]
+    def remove_listener(listener)
+      @listener_bus.remove(listener)
+      nil
+    end
+    alias removeListener remove_listener
+
+    # @return [Array<Object>] the registered listeners.
+    def listeners
+      @listener_bus.listeners
+    end
+
+    # Remove all listeners and stop the listener thread. @return [void]
+    def close
+      @listener_bus.close
+      nil
+    end
+
+    # Deliver a client-side event (e.g. QueryStartedEvent) to the listeners.
+    # @api private
+    def post_event(event)
+      @listener_bus.post(event)
     end
 
     # @return [Array<StreamingQuery>] the currently active queries.
@@ -348,4 +490,177 @@ module SparkConnect
       @session.client.execute_command(Proto::Command.new(streaming_query_manager_command: cmd)).streaming_manager_result
     end
   end
+
+  # Base class for streaming query listeners registered with
+  # {StreamingQueryManager#add_listener}. Override the callbacks you need; the
+  # camelCase PySpark names (`onQueryStarted`, ...) are honoured too.
+  class StreamingQueryListener
+    # @param event [QueryStartedEvent]
+    def on_query_started(event); end
+
+    # @param event [QueryProgressEvent]
+    def on_query_progress(event); end
+
+    # @param event [QueryIdleEvent]
+    def on_query_idle(event); end
+
+    # @param event [QueryTerminatedEvent]
+    def on_query_terminated(event); end
+
+    # A query started. `job_tags` is an Array of Strings.
+    QueryStartedEvent = Struct.new(:id, :run_id, :name, :timestamp, :job_tags, keyword_init: true)
+    # A micro-batch made progress. `progress` is the parsed progress Hash.
+    QueryProgressEvent = Struct.new(:progress, keyword_init: true)
+    # A query is idle, waiting for new data.
+    QueryIdleEvent = Struct.new(:id, :run_id, :timestamp, keyword_init: true)
+    # A query terminated, possibly with an exception message and error class.
+    QueryTerminatedEvent = Struct.new(:id, :run_id, :exception, :error_class_on_exception, keyword_init: true)
+
+    EVENT_TYPES = {
+      QUERY_PROGRESS_EVENT: :progress, QUERY_TERMINATED_EVENT: :terminated, QUERY_IDLE_EVENT: :idle,
+    }.freeze
+
+    CALLBACKS = {
+      QueryStartedEvent => %i[on_query_started onQueryStarted],
+      QueryProgressEvent => %i[on_query_progress onQueryProgress],
+      QueryIdleEvent => %i[on_query_idle onQueryIdle],
+      QueryTerminatedEvent => %i[on_query_terminated onQueryTerminated],
+    }.freeze
+
+    class << self
+      # Build an event from its server-side JSON.
+      # @api private
+      def event_from_json(kind, json)
+        j = JSON.parse(json)
+        case kind
+        when :started
+          QueryStartedEvent.new(id: j["id"], run_id: j["runId"], name: j["name"], timestamp: j["timestamp"],
+                                job_tags: Array(j["jobTags"]))
+        when :progress then QueryProgressEvent.new(progress: j["progress"])
+        when :idle then QueryIdleEvent.new(id: j["id"], run_id: j["runId"], timestamp: j["timestamp"])
+        when :terminated
+          QueryTerminatedEvent.new(id: j["id"], run_id: j["runId"], exception: j["exception"],
+                                   error_class_on_exception: j["errorClassOnException"])
+        end
+      end
+
+      # @api private - invoke the listener callback for `event`, if defined.
+      # A camelCase callback (PySpark's name) wins over the inherited no-op
+      # snake_case one of this base class.
+      def dispatch(listener, event)
+        snake, camel = CALLBACKS.fetch(event.class)
+        if listener.respond_to?(camel) &&
+           (!listener.respond_to?(snake) || listener.method(snake).owner == StreamingQueryListener)
+          listener.public_send(camel, event)
+        elsif listener.respond_to?(snake)
+          listener.public_send(snake, event)
+        end
+      end
+    end
+  end
+
+  # The client-side listener bus behind {StreamingQueryManager#add_listener}:
+  # registers one server-side listener for the session and dispatches the
+  # events it streams back to every registered listener on a background thread.
+  # @api private
+  class StreamingQueryListenerBus
+    Proto = SparkConnect::Proto
+
+    def initialize(session)
+      @session = session
+      @listeners = []
+      @lock = Mutex.new
+      @thread = nil
+    end
+
+    # @return [Array<Object>] a snapshot of the registered listeners.
+    def listeners
+      @lock.synchronize { @listeners.dup }
+    end
+
+    # Register `listener`; the first one starts the server-side listener.
+    # @return [void]
+    def add(listener)
+      @lock.synchronize do
+        return if @listeners.include?(listener)
+
+        @listeners << listener
+        start_server_listener if @listeners.size == 1
+      end
+    end
+
+    # Unregister `listener`; the last one stops the server-side listener.
+    # @return [void]
+    def remove(listener)
+      thread = nil
+      @lock.synchronize do
+        return unless @listeners.include?(listener)
+
+        if @listeners.size == 1
+          bus_command(remove_listener_bus_listener: true)
+          thread = @thread
+          @thread = nil
+        end
+        @listeners.delete(listener)
+      end
+      # Join outside the lock: the handler thread takes it to dispatch events.
+      thread&.join(30)
+    end
+
+    # Unregister every listener. @return [void]
+    def close
+      listeners.each { |l| remove(l) }
+    end
+
+    # Deliver `event` to every listener (exceptions are reported as warnings).
+    # @return [void]
+    def post(event)
+      listeners.each do |listener|
+        StreamingQueryListener.dispatch(listener, event)
+      rescue StandardError => e
+        warn("spark-connect: streaming listener #{listener.class} raised #{e.class}: #{e.message}")
+      end
+    end
+
+    private
+
+    def bus_command(**kw)
+      cmd = Proto::Command.new(streaming_query_listener_bus_command: Proto::StreamingQueryListenerBusCommand.new(**kw))
+      @session.client.execute_command(cmd)
+    end
+
+    # Must be called with @lock held.
+    def start_server_listener
+      ready = Queue.new
+      cmd = Proto::Command.new(
+        streaming_query_listener_bus_command: Proto::StreamingQueryListenerBusCommand.new(add_listener_bus_listener: true)
+      )
+      @thread = Thread.new do
+        @session.client.stream_command(cmd) do |resp|
+          result = resp.streaming_query_listener_events_result
+          next unless result
+
+          ready << :added if result.listener_bus_listener_added
+          result.events.each do |e|
+            kind = StreamingQueryListener::EVENT_TYPES[e.event_type]
+            post(StreamingQueryListener.event_from_json(kind, e.event_json)) if kind
+          end
+        end
+      rescue StandardError => e
+        ready << e
+        warn("spark-connect: streaming listener bus stopped: #{e.class}: #{e.message}; re-add the listeners")
+        @lock.synchronize { @listeners.clear }
+      ensure
+        ready << :finished
+      end
+      status = ready.pop
+      return if status == :added
+
+      @thread = nil
+      @listeners.clear
+      raise status if status.is_a?(Exception)
+
+      raise SparkConnectError, "The server did not register the streaming query listener"
+    end
+  end
 end

data/lib/spark_connect/table_arg.rb

@@ -0,0 +1,73 @@
+# frozen_string_literal: true
+
+module SparkConnect
+  # A {DataFrame} passed as a `TABLE` argument to a table-valued function,
+  # created with {DataFrame#as_table}. It is a {Column} (a table-argument
+  # subquery expression), so it can be passed wherever a function argument is
+  # accepted, e.g. `F.call_function("my_tvf", df.as_table.partition_by("k"))`.
+  #
+  # Mirrors PySpark's `TableArg`: optionally partition the rows
+  # ({#partition_by} or {#with_single_partition}) and then order them within
+  # each partition ({#order_by}).
+  class TableArg < Column
+    Proto = SparkConnect::Proto
+
+    # @param df [DataFrame]
+    # @api private
+    def initialize(df, partition_spec: [], order_spec: [], with_single_partition: nil)
+      @df = df
+      @partition_spec = partition_spec
+      @order_spec = order_spec
+      @with_single_partition = with_single_partition
+      options = nil
+      unless partition_spec.empty? && order_spec.empty? && with_single_partition.nil?
+        options = Proto::SubqueryExpression::TableArgOptions.new(partition_spec: partition_spec, order_spec: order_spec)
+        options.with_single_partition = with_single_partition unless with_single_partition.nil?
+      end
+      super(df.subquery_expression(:SUBQUERY_TYPE_TABLE_ARG, table_arg_options: options).to_expr)
+    end
+
+    # Partition the table argument's rows by these columns.
+    # @param cols [Array<Column, String>]
+    # @return [TableArg]
+    def partition_by(*cols)
+      raise IllegalArgumentError, "Cannot call partition_by after partition_by or with_single_partition" if partitioned?
+
+      exprs = cols.flatten.map { |c| Functions._col(c).to_expr }
+      TableArg.new(@df, partition_spec: @partition_spec + exprs, order_spec: @order_spec,
+                        with_single_partition: @with_single_partition)
+    end
+    alias partitionBy partition_by
+
+    # Order the rows within each partition (after {#partition_by} or
+    # {#with_single_partition}).
+    # @param cols [Array<Column, String>]
+    # @return [TableArg]
+    def order_by(*cols)
+      raise IllegalArgumentError, "Call partition_by or with_single_partition before order_by" unless partitioned?
+
+      orders = cols.flatten.map do |c|
+        expr = Functions._col(c).to_expr
+        expr.expr_type == :sort_order ? expr.sort_order : Functions._col(c).asc.to_expr.sort_order
+      end
+      TableArg.new(@df, partition_spec: @partition_spec, order_spec: @order_spec + orders,
+                        with_single_partition: @with_single_partition)
+    end
+    alias orderBy order_by
+
+    # Pass all rows to a single partition.
+    # @return [TableArg]
+    def with_single_partition
+      raise IllegalArgumentError, "Cannot call with_single_partition after partition_by or with_single_partition" if partitioned?
+
+      TableArg.new(@df, partition_spec: @partition_spec, order_spec: @order_spec, with_single_partition: true)
+    end
+    alias withSinglePartition with_single_partition
+
+    private
+
+    def partitioned?
+      !@partition_spec.empty? || @with_single_partition == true
+    end
+  end
+end