spark-connect 0.2.0 → 0.4.0

data/lib/spark_connect/proto/spark/connect/types_pb.rb

@@ -5,7 +5,7 @@
 require 'google/protobuf'
 
 
-descriptor_data = "\n\x19spark/connect/types.proto\x12\rspark.connect\"\xbc\x1c\n\x08\x44\x61taType\x12,\n\x04null\x18\x01 \x01(\x0b\x32\x1c.spark.connect.DataType.NULLH\x00\x12\x30\n\x06\x62inary\x18\x02 \x01(\x0b\x32\x1e.spark.connect.DataType.BinaryH\x00\x12\x32\n\x07\x62oolean\x18\x03 \x01(\x0b\x32\x1f.spark.connect.DataType.BooleanH\x00\x12,\n\x04\x62yte\x18\x04 \x01(\x0b\x32\x1c.spark.connect.DataType.ByteH\x00\x12.\n\x05short\x18\x05 \x01(\x0b\x32\x1d.spark.connect.DataType.ShortH\x00\x12\x32\n\x07integer\x18\x06 \x01(\x0b\x32\x1f.spark.connect.DataType.IntegerH\x00\x12,\n\x04long\x18\x07 \x01(\x0b\x32\x1c.spark.connect.DataType.LongH\x00\x12.\n\x05\x66loat\x18\x08 \x01(\x0b\x32\x1d.spark.connect.DataType.FloatH\x00\x12\x30\n\x06\x64ouble\x18\t \x01(\x0b\x32\x1e.spark.connect.DataType.DoubleH\x00\x12\x32\n\x07\x64\x65\x63imal\x18\n \x01(\x0b\x32\x1f.spark.connect.DataType.DecimalH\x00\x12\x30\n\x06string\x18\x0b \x01(\x0b\x32\x1e.spark.connect.DataType.StringH\x00\x12,\n\x04\x63har\x18\x0c \x01(\x0b\x32\x1c.spark.connect.DataType.CharH\x00\x12\x33\n\x08var_char\x18\r \x01(\x0b\x32\x1f.spark.connect.DataType.VarCharH\x00\x12,\n\x04\x64\x61te\x18\x0e \x01(\x0b\x32\x1c.spark.connect.DataType.DateH\x00\x12\x36\n\ttimestamp\x18\x0f \x01(\x0b\x32!.spark.connect.DataType.TimestampH\x00\x12=\n\rtimestamp_ntz\x18\x10 \x01(\x0b\x32$.spark.connect.DataType.TimestampNTZH\x00\x12\x45\n\x11\x63\x61lendar_interval\x18\x11 \x01(\x0b\x32(.spark.connect.DataType.CalendarIntervalH\x00\x12H\n\x13year_month_interval\x18\x12 \x01(\x0b\x32).spark.connect.DataType.YearMonthIntervalH\x00\x12\x44\n\x11\x64\x61y_time_interval\x18\x13 \x01(\x0b\x32\'.spark.connect.DataType.DayTimeIntervalH\x00\x12.\n\x05\x61rray\x18\x14 \x01(\x0b\x32\x1d.spark.connect.DataType.ArrayH\x00\x12\x30\n\x06struct\x18\x15 \x01(\x0b\x32\x1e.spark.connect.DataType.StructH\x00\x12*\n\x03map\x18\x16 \x01(\x0b\x32\x1b.spark.connect.DataType.MapH\x00\x12\x32\n\x07variant\x18\x19 \x01(\x0b\x32\x1f.spark.connect.DataType.VariantH\x00\x12*\n\x03udt\x18\x17 \x01(\x0b\x32\x1b.spark.connect.DataType.UDTH\x00\x12\x34\n\x08geometry\x18\x1a \x01(\x0b\x32 .spark.connect.DataType.GeometryH\x00\x12\x36\n\tgeography\x18\x1b \x01(\x0b\x32!.spark.connect.DataType.GeographyH\x00\x12\x34\n\x08unparsed\x18\x18 \x01(\x0b\x32 .spark.connect.DataType.UnparsedH\x00\x12,\n\x04time\x18\x1c \x01(\x0b\x32\x1c.spark.connect.DataType.TimeH\x00\x1a+\n\x07\x42oolean\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a(\n\x04\x42yte\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a)\n\x05Short\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a+\n\x07Integer\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a(\n\x04Long\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a)\n\x05\x46loat\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a*\n\x06\x44ouble\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a=\n\x06String\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x12\x11\n\tcollation\x18\x02 \x01(\t\x1a*\n\x06\x42inary\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a(\n\x04NULL\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a-\n\tTimestamp\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a(\n\x04\x44\x61te\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a\x30\n\x0cTimestampNTZ\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1aN\n\x04Time\x12\x16\n\tprecision\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x02 \x01(\rB\x0c\n\n_precision\x1a\x34\n\x10\x43\x61lendarInterval\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a\x85\x01\n\x11YearMonthInterval\x12\x18\n\x0bstart_field\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x16\n\tend_field\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x03 \x01(\rB\x0e\n\x0c_start_fieldB\x0c\n\n_end_field\x1a\x83\x01\n\x0f\x44\x61yTimeInterval\x12\x18\n\x0bstart_field\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x16\n\tend_field\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x03 \x01(\rB\x0e\n\x0c_start_fieldB\x0c\n\n_end_field\x1a\x38\n\x04\x43har\x12\x0e\n\x06length\x18\x01 \x01(\x05\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1a;\n\x07VarChar\x12\x0e\n\x06length\x18\x01 \x01(\x05\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1ao\n\x07\x44\x65\x63imal\x12\x12\n\x05scale\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x16\n\tprecision\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x03 \x01(\rB\x08\n\x06_scaleB\x0c\n\n_precision\x1a}\n\x0bStructField\x12\x0c\n\x04name\x18\x01 \x01(\t\x12*\n\tdata_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataType\x12\x10\n\x08nullable\x18\x03 \x01(\x08\x12\x15\n\x08metadata\x18\x04 \x01(\tH\x00\x88\x01\x01\x42\x0b\n\t_metadata\x1a_\n\x06Struct\x12\x33\n\x06\x66ields\x18\x01 \x03(\x0b\x32#.spark.connect.DataType.StructField\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1ao\n\x05\x41rray\x12-\n\x0c\x65lement_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataType\x12\x15\n\rcontains_null\x18\x02 \x01(\x08\x12 \n\x18type_variation_reference\x18\x03 \x01(\r\x1a\x9c\x01\n\x03Map\x12)\n\x08key_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataType\x12+\n\nvalue_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataType\x12\x1b\n\x13value_contains_null\x18\x03 \x01(\x08\x12 \n\x18type_variation_reference\x18\x04 \x01(\r\x1a:\n\x08Geometry\x12\x0c\n\x04srid\x18\x01 \x01(\x05\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1a;\n\tGeography\x12\x0c\n\x04srid\x18\x01 \x01(\x05\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1a+\n\x07Variant\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a\xe4\x01\n\x03UDT\x12\x0c\n\x04type\x18\x01 \x01(\t\x12\x16\n\tjvm_class\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0cpython_class\x18\x03 \x01(\tH\x01\x88\x01\x01\x12$\n\x17serialized_python_class\x18\x04 \x01(\tH\x02\x88\x01\x01\x12.\n\x08sql_type\x18\x05 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x03\x88\x01\x01\x42\x0c\n\n_jvm_classB\x0f\n\r_python_classB\x1a\n\x18_serialized_python_classB\x0b\n\t_sql_type\x1a$\n\x08Unparsed\x12\x18\n\x10\x64\x61ta_type_string\x18\x01 \x01(\tB\x06\n\x04kindB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3"
+descriptor_data = "\n\x19spark/connect/types.proto\x12\rspark.connect\"\x8a\x1f\n\x08\x44\x61taType\x12,\n\x04null\x18\x01 \x01(\x0b\x32\x1c.spark.connect.DataType.NULLH\x00\x12\x30\n\x06\x62inary\x18\x02 \x01(\x0b\x32\x1e.spark.connect.DataType.BinaryH\x00\x12\x32\n\x07\x62oolean\x18\x03 \x01(\x0b\x32\x1f.spark.connect.DataType.BooleanH\x00\x12,\n\x04\x62yte\x18\x04 \x01(\x0b\x32\x1c.spark.connect.DataType.ByteH\x00\x12.\n\x05short\x18\x05 \x01(\x0b\x32\x1d.spark.connect.DataType.ShortH\x00\x12\x32\n\x07integer\x18\x06 \x01(\x0b\x32\x1f.spark.connect.DataType.IntegerH\x00\x12,\n\x04long\x18\x07 \x01(\x0b\x32\x1c.spark.connect.DataType.LongH\x00\x12.\n\x05\x66loat\x18\x08 \x01(\x0b\x32\x1d.spark.connect.DataType.FloatH\x00\x12\x30\n\x06\x64ouble\x18\t \x01(\x0b\x32\x1e.spark.connect.DataType.DoubleH\x00\x12\x32\n\x07\x64\x65\x63imal\x18\n \x01(\x0b\x32\x1f.spark.connect.DataType.DecimalH\x00\x12\x30\n\x06string\x18\x0b \x01(\x0b\x32\x1e.spark.connect.DataType.StringH\x00\x12,\n\x04\x63har\x18\x0c \x01(\x0b\x32\x1c.spark.connect.DataType.CharH\x00\x12\x33\n\x08var_char\x18\r \x01(\x0b\x32\x1f.spark.connect.DataType.VarCharH\x00\x12,\n\x04\x64\x61te\x18\x0e \x01(\x0b\x32\x1c.spark.connect.DataType.DateH\x00\x12\x36\n\ttimestamp\x18\x0f \x01(\x0b\x32!.spark.connect.DataType.TimestampH\x00\x12=\n\rtimestamp_ntz\x18\x10 \x01(\x0b\x32$.spark.connect.DataType.TimestampNTZH\x00\x12\x45\n\x11\x63\x61lendar_interval\x18\x11 \x01(\x0b\x32(.spark.connect.DataType.CalendarIntervalH\x00\x12H\n\x13year_month_interval\x18\x12 \x01(\x0b\x32).spark.connect.DataType.YearMonthIntervalH\x00\x12\x44\n\x11\x64\x61y_time_interval\x18\x13 \x01(\x0b\x32\'.spark.connect.DataType.DayTimeIntervalH\x00\x12.\n\x05\x61rray\x18\x14 \x01(\x0b\x32\x1d.spark.connect.DataType.ArrayH\x00\x12\x30\n\x06struct\x18\x15 \x01(\x0b\x32\x1e.spark.connect.DataType.StructH\x00\x12*\n\x03map\x18\x16 \x01(\x0b\x32\x1b.spark.connect.DataType.MapH\x00\x12\x32\n\x07variant\x18\x19 \x01(\x0b\x32\x1f.spark.connect.DataType.VariantH\x00\x12*\n\x03udt\x18\x17 \x01(\x0b\x32\x1b.spark.connect.DataType.UDTH\x00\x12\x34\n\x08geometry\x18\x1a \x01(\x0b\x32 .spark.connect.DataType.GeometryH\x00\x12\x36\n\tgeography\x18\x1b \x01(\x0b\x32!.spark.connect.DataType.GeographyH\x00\x12\x34\n\x08unparsed\x18\x18 \x01(\x0b\x32 .spark.connect.DataType.UnparsedH\x00\x12,\n\x04time\x18\x1c \x01(\x0b\x32\x1c.spark.connect.DataType.TimeH\x00\x12H\n\x13timestamp_ntz_nanos\x18\x1d \x01(\x0b\x32).spark.connect.DataType.TimestampNTZNanosH\x00\x12H\n\x13timestamp_ltz_nanos\x18\x1e \x01(\x0b\x32).spark.connect.DataType.TimestampLTZNanosH\x00\x1a+\n\x07\x42oolean\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a(\n\x04\x42yte\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a)\n\x05Short\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a+\n\x07Integer\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a(\n\x04Long\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a)\n\x05\x46loat\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a*\n\x06\x44ouble\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a=\n\x06String\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x12\x11\n\tcollation\x18\x02 \x01(\t\x1a*\n\x06\x42inary\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a(\n\x04NULL\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a-\n\tTimestamp\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a(\n\x04\x44\x61te\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a\x30\n\x0cTimestampNTZ\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1aN\n\x04Time\x12\x16\n\tprecision\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x02 \x01(\rB\x0c\n\n_precision\x1a[\n\x11TimestampNTZNanos\x12\x16\n\tprecision\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x02 \x01(\rB\x0c\n\n_precision\x1a[\n\x11TimestampLTZNanos\x12\x16\n\tprecision\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x02 \x01(\rB\x0c\n\n_precision\x1a\x34\n\x10\x43\x61lendarInterval\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a\x85\x01\n\x11YearMonthInterval\x12\x18\n\x0bstart_field\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x16\n\tend_field\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x03 \x01(\rB\x0e\n\x0c_start_fieldB\x0c\n\n_end_field\x1a\x83\x01\n\x0f\x44\x61yTimeInterval\x12\x18\n\x0bstart_field\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x16\n\tend_field\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x03 \x01(\rB\x0e\n\x0c_start_fieldB\x0c\n\n_end_field\x1a\x38\n\x04\x43har\x12\x0e\n\x06length\x18\x01 \x01(\x05\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1a;\n\x07VarChar\x12\x0e\n\x06length\x18\x01 \x01(\x05\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1ao\n\x07\x44\x65\x63imal\x12\x12\n\x05scale\x18\x01 \x01(\x05H\x00\x88\x01\x01\x12\x16\n\tprecision\x18\x02 \x01(\x05H\x01\x88\x01\x01\x12 \n\x18type_variation_reference\x18\x03 \x01(\rB\x08\n\x06_scaleB\x0c\n\n_precision\x1a}\n\x0bStructField\x12\x0c\n\x04name\x18\x01 \x01(\t\x12*\n\tdata_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataType\x12\x10\n\x08nullable\x18\x03 \x01(\x08\x12\x15\n\x08metadata\x18\x04 \x01(\tH\x00\x88\x01\x01\x42\x0b\n\t_metadata\x1a_\n\x06Struct\x12\x33\n\x06\x66ields\x18\x01 \x03(\x0b\x32#.spark.connect.DataType.StructField\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1ao\n\x05\x41rray\x12-\n\x0c\x65lement_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataType\x12\x15\n\rcontains_null\x18\x02 \x01(\x08\x12 \n\x18type_variation_reference\x18\x03 \x01(\r\x1a\x9c\x01\n\x03Map\x12)\n\x08key_type\x18\x01 \x01(\x0b\x32\x17.spark.connect.DataType\x12+\n\nvalue_type\x18\x02 \x01(\x0b\x32\x17.spark.connect.DataType\x12\x1b\n\x13value_contains_null\x18\x03 \x01(\x08\x12 \n\x18type_variation_reference\x18\x04 \x01(\r\x1a:\n\x08Geometry\x12\x0c\n\x04srid\x18\x01 \x01(\x05\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1a;\n\tGeography\x12\x0c\n\x04srid\x18\x01 \x01(\x05\x12 \n\x18type_variation_reference\x18\x02 \x01(\r\x1a+\n\x07Variant\x12 \n\x18type_variation_reference\x18\x01 \x01(\r\x1a\xe4\x01\n\x03UDT\x12\x0c\n\x04type\x18\x01 \x01(\t\x12\x16\n\tjvm_class\x18\x02 \x01(\tH\x00\x88\x01\x01\x12\x19\n\x0cpython_class\x18\x03 \x01(\tH\x01\x88\x01\x01\x12$\n\x17serialized_python_class\x18\x04 \x01(\tH\x02\x88\x01\x01\x12.\n\x08sql_type\x18\x05 \x01(\x0b\x32\x17.spark.connect.DataTypeH\x03\x88\x01\x01\x42\x0c\n\n_jvm_classB\x0f\n\r_python_classB\x1a\n\x18_serialized_python_classB\x0b\n\t_sql_type\x1a$\n\x08Unparsed\x12\x18\n\x10\x64\x61ta_type_string\x18\x01 \x01(\tB\x06\n\x04kindB6\n\x1eorg.apache.spark.connect.protoP\x01Z\x12internal/generatedb\x06proto3"
 
 pool = ::Google::Protobuf::DescriptorPool.generated_pool
 pool.add_serialized_file(descriptor_data)
@@ -27,6 +27,8 @@ module Spark
     DataType::Date = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("spark.connect.DataType.Date").msgclass
     DataType::TimestampNTZ = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("spark.connect.DataType.TimestampNTZ").msgclass
     DataType::Time = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("spark.connect.DataType.Time").msgclass
+    DataType::TimestampNTZNanos = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("spark.connect.DataType.TimestampNTZNanos").msgclass
+    DataType::TimestampLTZNanos = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("spark.connect.DataType.TimestampLTZNanos").msgclass
     DataType::CalendarInterval = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("spark.connect.DataType.CalendarInterval").msgclass
     DataType::YearMonthInterval = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("spark.connect.DataType.YearMonthInterval").msgclass
     DataType::DayTimeInterval = ::Google::Protobuf::DescriptorPool.generated_pool.lookup("spark.connect.DataType.DayTimeInterval").msgclass

data/lib/spark_connect/reader.rb

@@ -8,6 +8,7 @@ module SparkConnect
   #   spark.read.format("csv").option("header", true).load("data.csv")
   #   spark.read.json("events.json")
   #   spark.read.table("my_table")
+  #   spark.read.json(lines_df)  # parse a DataFrame of JSON strings
   class DataFrameReader
     Proto = SparkConnect::Proto
 
@@ -29,21 +30,25 @@ module SparkConnect
     # Set the input schema (a {Types::StructType} or DDL string).
     # @return [self]
     def schema(schema)
-      @schema = schema.is_a?(Types::StructType) ? schema.simple_string : schema.to_s
+      unless schema.is_a?(Types::StructType) || schema.is_a?(String)
+        raise IllegalArgumentError, "schema must be a StructType or DDL String, got #{schema.class}"
+      end
+
+      @schema = schema
       self
     end
 
-    # Set a single read option.
+    # Set a single read option. A `nil` value is ignored.
     # @return [self]
     def option(key, value)
-      @options[key.to_s] = value.to_s
+      @options[key.to_s] = value.to_s unless value.nil?
       self
     end
 
     # Set multiple read options.
     # @return [self]
     def options(opts)
-      opts.each { |k, v| @options[k.to_s] = v.to_s }
+      opts.each { |k, v| option(k, v) }
       self
     end
 
@@ -52,10 +57,7 @@ module SparkConnect
     # @param paths [Array<String>]
     # @return [DataFrame]
     def load(*paths)
-      ds = Proto::Read::DataSource.new(options: @options, paths: paths.flatten.map(&:to_s))
-      ds.format = @format if @format
-      ds.schema = @schema if @schema
-      read_relation(data_source: ds)
+      read_relation(data_source: data_source(paths.flatten.map(&:to_s)))
     end
 
     # Read a registered table or view.
@@ -67,10 +69,31 @@ module SparkConnect
       read_relation(named_table: nt)
     end
 
-    # @return [DataFrame] CSV at `paths`.
-    def csv(*paths) = format("csv").load(*paths)
-    # @return [DataFrame] JSON at `paths`.
-    def json(*paths) = format("json").load(*paths)
+    # The row-level changes (Change Data Capture) of a table as a batch
+    # DataFrame (Spark 4.2+; the table's catalog must support change logs).
+    # Configure the range with options such as `startingVersion`,
+    # `endingVersion`, `startingTimestamp`, `endingTimestamp`,
+    # `deduplicationMode`, and `computeUpdates`.
+    #
+    # @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)
+      DataFrame.new(@session, PlanBuilder.relation(@session, relation_changes: rc))
+    end
+
+    # CSV at `paths`, or parse a single-string-column {DataFrame} of CSV lines.
+    # @return [DataFrame]
+    def csv(*paths) = text_source("csv", :PARSE_FORMAT_CSV, paths)
+    # JSON at `paths`, or parse a single-string-column {DataFrame} of JSON documents.
+    # @return [DataFrame]
+    def json(*paths) = text_source("json", :PARSE_FORMAT_JSON, paths)
+    # XML at `paths` (set `option("rowTag", ...)`), or parse a single-string-column
+    # {DataFrame} of XML records (Spark 4.2+ for the DataFrame form).
+    # @return [DataFrame]
+    def xml(*paths) = text_source("xml", :PARSE_FORMAT_XML, paths)
     # @return [DataFrame] Parquet at `paths`.
     def parquet(*paths) = format("parquet").load(*paths)
     # @return [DataFrame] ORC at `paths`.
@@ -78,19 +101,77 @@ module SparkConnect
     # @return [DataFrame] text at `paths` (one `value` column per line).
     def text(*paths) = format("text").load(*paths)
 
-    # Read from a JDBC source.
+    JDBC_PARTITION_KEYS = %i[column lower_bound upper_bound num_partitions predicates].freeze
+
+    # Read from a JDBC source, optionally partitioned on a numeric column
+    # (`column:` with `lower_bound:`, `upper_bound:`, `num_partitions:`) or by a
+    # list of WHERE-clause `predicates:` (one partition each).
     #
+    # @example
+    #   spark.read.jdbc(url, "people", "user" => "u", "password" => "p")
+    #   spark.read.jdbc(url, "people", { "user" => "u" }, column: "id", lower_bound: 0, upper_bound: 100, num_partitions: 4)
     # @param url [String] the JDBC URL.
     # @param table [String] the table name (or subquery).
-    # @param properties [Hash] connection properties (`user`, `password`, ...).
+    # @param properties [Hash] connection properties (`user`, `password`, ...);
+    #   String-keyed keyword arguments are merged into it.
     # @return [DataFrame]
-    def jdbc(url, table, properties = {})
-      opts = { "url" => url, "dbtable" => table }.merge(properties.transform_keys(&:to_s))
-      format("jdbc").options(opts).load
+    def jdbc(url, table, properties = {}, **opts)
+      # Symbol keys are partitioning options; anything else is a connection property.
+      partitioning, extra = opts.partition { |k, _| JDBC_PARTITION_KEYS.include?(k) }.map(&:to_h)
+      properties = properties.merge(extra)
+      column, lower_bound, upper_bound, num_partitions, predicates = partitioning.values_at(*JDBC_PARTITION_KEYS)
+      format("jdbc")
+      unless column.nil?
+        if lower_bound.nil? || upper_bound.nil? || num_partitions.nil?
+          raise IllegalArgumentError, "lower_bound, upper_bound, and num_partitions are required with column"
+        end
+
+        options("partitionColumn" => column, "lowerBound" => lower_bound, "upperBound" => upper_bound,
+                "numPartitions" => num_partitions)
+      end
+      options(properties.transform_keys(&:to_s))
+      options("url" => url, "dbtable" => table)
+      return load if predicates.nil? || !column.nil?
+
+      read_relation(data_source: data_source([], predicates: predicates.map(&:to_s)))
     end
 
     private
 
+    def data_source(paths, predicates: [])
+      ds = Proto::Read::DataSource.new(options: @options, paths: paths, predicates: predicates)
+      ds.format = @format if @format
+      ds.schema = schema_string if @schema
+      ds
+    end
+
+    # Like PySpark, a StructType is sent as its JSON form (which the server
+    # parses losslessly); a String is sent as given (DDL or JSON).
+    def schema_string
+      @schema.is_a?(Types::StructType) ? @schema.json : @schema
+    end
+
+    def text_source(fmt, parse_format, paths)
+      return parse(paths.first, parse_format) if paths.size == 1 && paths.first.is_a?(DataFrame)
+
+      format(fmt).load(*paths)
+    end
+
+    def parse(df, parse_format)
+      parse = Proto::Parse.new(input: df.relation, format: parse_format, options: @options)
+      parse.schema = parse_schema.to_proto if @schema
+      DataFrame.new(@session, PlanBuilder.relation(@session, parse: parse))
+    end
+
+    # The Parse relation carries a typed schema; resolve a DDL string through
+    # the server's parser.
+    def parse_schema
+      return @schema if @schema.is_a?(Types::StructType)
+
+      parsed = @session.client.analyze(ddl_parse: Proto::AnalyzePlanRequest::DDLParse.new(ddl_string: @schema)).ddl_parse.parsed
+      Types.from_proto(parsed)
+    end
+
     def read_relation(**read_kw)
       DataFrame.new(@session, PlanBuilder.relation(@session, read: Proto::Read.new(**read_kw)))
     end

data/lib/spark_connect/session.rb

@@ -25,26 +25,99 @@ module SparkConnect
     def initialize(client)
       @client = client
       @plan_id = -1
+      @plan_id_lock = Mutex.new
       @conf = RuntimeConfig.new(client)
+      @subquery_relations = {}
+      @release_session_on_close = true
     end
 
+    # Whether {#stop} releases the server-side session.
+    # @api private
+    attr_accessor :release_session_on_close
+
+    @default_session = nil
+    @session_lock = Mutex.new
+    ACTIVE_SESSION_KEY = :spark_connect_active_session
+
     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}.
+      # The active session of the current thread, else the process-wide
+      # default session (the first session built by {Builder#get_or_create} /
+      # {Builder#create}); stopped sessions are skipped.
+      #
+      # Unlike PySpark's `SparkSession.active()`, this returns nil (instead of
+      # raising) when there is no session, which is the idiomatic Ruby shape for
+      # an optional lookup.
+      #
       # @return [SparkSession, nil]
-      attr_accessor :active
+      def active
+        get_active_session || default_session
+      end
+
+      # Make `session` the active session of the current thread (and the
+      # process-wide default when none is set). `nil` clears the active session.
+      # @param session [SparkSession, nil]
+      def active=(session)
+        Thread.current[ACTIVE_SESSION_KEY] = session
+        @session_lock.synchronize { @default_session ||= session } if session
+      end
+
+      # @return [SparkSession, nil] the active session of the current thread, if not stopped.
+      def get_active_session
+        s = Thread.current[ACTIVE_SESSION_KEY]
+        s && !s.stopped? ? s : nil
+      end
+      alias getActiveSession get_active_session
 
+      # @return [SparkSession, nil] the process-wide default session, if not stopped.
+      def default_session
+        s = @session_lock.synchronize { @default_session }
+        s && !s.stopped? ? s : nil
+      end
+
+      # Clear the process-wide default session (it is not stopped). @return [void]
+      def clear_default_session
+        @session_lock.synchronize { @default_session = nil }
+        nil
+      end
+
+      # Clear the current thread's active session (it is not stopped). @return [void]
+      def clear_active_session
+        Thread.current[ACTIVE_SESSION_KEY] = nil
+      end
+
+      # Forget `session` wherever it is the default/active one.
       # @api private
+      def forget(session)
+        @session_lock.synchronize { @default_session = nil if @default_session.equal?(session) }
+        Thread.current[ACTIVE_SESSION_KEY] = nil if Thread.current[ACTIVE_SESSION_KEY].equal?(session)
+      end
     end
 
     # Allocate the next unique plan id. Used by {PlanBuilder.relation}.
     # @return [Integer]
     def next_plan_id
-      @plan_id += 1
+      @plan_id_lock.synchronize { @plan_id += 1 }
+    end
+
+    # @api private - remember a DataFrame's relation so a subquery expression
+    #   that refers to it by plan id can be expanded into a `WithRelations`.
+    # @param relation [Spark::Connect::Relation]
+    # @return [Integer] the relation's plan id.
+    def register_subquery_relation(relation)
+      plan_id = relation.common.plan_id
+      @plan_id_lock.synchronize { @subquery_relations[plan_id] = relation }
+      plan_id
+    end
+
+    # @api private
+    # @return [Spark::Connect::Relation, nil]
+    def subquery_relation(plan_id)
+      @plan_id_lock.synchronize { @subquery_relations[plan_id] }
     end
 
     # @return [String] the client session id (UUID).
@@ -105,7 +178,13 @@ module SparkConnect
 
     # @return [StreamingQueryManager] the manager for this session's streaming queries.
     def streams
-      StreamingQueryManager.new(self)
+      @streams ||= StreamingQueryManager.new(self)
+    end
+
+    # @return [TableValuedFunction] built-in table-valued functions
+    #   (`spark.tvf.explode(...)`, `spark.tvf.range(...)`, ...).
+    def tvf
+      TableValuedFunction.new(self)
     end
 
     # Create a new Spark Declarative Pipeline (a dataflow graph) in this session.
@@ -134,6 +213,15 @@ module SparkConnect
     alias create_dataframe create_data_frame
     alias createDataFrame create_data_frame
 
+    # An empty {DataFrame} with the given schema.
+    #
+    # @param schema [Types::StructType, String] a schema or DDL string.
+    # @return [DataFrame]
+    def empty_data_frame(schema)
+      create_data_frame([], schema)
+    end
+    alias emptyDataFrame empty_data_frame
+
     # @return [RuntimeConfig] runtime configuration facade.
     attr_reader :conf
 
@@ -147,19 +235,60 @@ module SparkConnect
       @client.analyze(spark_version: Proto::AnalyzePlanRequest::SparkVersion.new).spark_version.version
     end
 
-    # Make this the active/default session.
+    # Make this the active (and, if unset, the 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).
+    # Start a brand-new session against the same endpoint: a fresh server-side
+    # session with no configuration, temporary views, or functions carried over
+    # (Spark Connect `newSession` semantics, like PySpark 4.3). Configurations
+    # set through the builder are not reapplied.
     # @return [SparkSession]
     def new_session
-      SparkSession.new(SparkConnectClient.new(@client.channel_builder))
+      session = SparkSession.new(@client.new_session_client)
+      session.release_session_on_close = true
+      session
+    end
+    alias newSession new_session
+
+    # Clone this session on the server: the new session starts with a copy of
+    # this session's SQL configurations, temporary views, registered functions,
+    # and catalog state, and is isolated from it afterwards (Spark 4.1+).
+    #
+    # @param new_session_id [String, nil] a UUID for the clone (generated when omitted).
+    # @return [SparkSession]
+    def clone_session(new_session_id = nil)
+      session = SparkSession.new(@client.clone_session(new_session_id))
+      session.release_session_on_close = true
+      session
     end
+    alias cloneSession clone_session
+
+    # Register a callable invoked with execution progress of every query run by
+    # this session. It is called with the keywords `stages:` (an Array of
+    # {StageInfo}), `inflight_tasks:`, `operation_id:`, and `done:`.
+    #
+    # @param handler [#call] a Proc/lambda or any object responding to `call`.
+    # @return [void]
+    def register_progress_handler(handler = nil, &block)
+      @client.register_progress_handler(handler || block)
+    end
+    alias registerProgressHandler register_progress_handler
+
+    # Remove a handler added with {#register_progress_handler}. @return [void]
+    def remove_progress_handler(handler)
+      @client.remove_progress_handler(handler)
+    end
+    alias removeProgressHandler remove_progress_handler
+
+    # Remove all progress handlers. @return [void]
+    def clear_progress_handlers
+      @client.clear_progress_handlers
+    end
+    alias clearProgressHandlers clear_progress_handlers
 
     # Interrupt all operations running in this session.
     # @return [Array<String>] the ids of the interrupted operations.
@@ -200,14 +329,24 @@ module SparkConnect
       @client.clear_tags
     end
 
-    # Release the server-side session and stop the client.
+    # Release the server-side session and stop the client. Idempotent.
     # @return [void]
     def stop
-      @client.release_session
-      SparkSession.active = nil if SparkSession.active.equal?(self)
+      return nil if stopped?
+
+      @streams&.close
+      @client.release_session if @release_session_on_close
+      @client.close
+      SparkSession.forget(self)
       nil
     end
 
+    # @return [Boolean] whether {#stop} has been called.
+    def stopped?
+      @client.respond_to?(:closed?) && @client.closed?
+    end
+    alias is_stopped stopped?
+
     # @api private
     def create_data_frame_from_relation(relation)
       DataFrame.new(self, relation)
@@ -264,9 +403,14 @@ module SparkConnect
   # Fluent builder for {SparkSession}. Returned by {SparkSession.builder}.
   class SparkSession
     class Builder
+      # Options that only make sense when connecting and are never applied as
+      # session configurations.
+      CONNECT_ONLY_OPTIONS = %w[spark.remote spark.master].freeze
+
       def initialize
         @options = {}
         @remote = nil
+        @channel_builder = nil
       end
 
       # Set the connection string (`sc://...`).
@@ -276,42 +420,100 @@ module SparkConnect
         self
       end
 
+      # Use a custom {ChannelBuilder} (e.g. a subclass that customizes the gRPC
+      # channel) instead of a connection string.
+      # @param channel_builder [ChannelBuilder]
+      # @return [self]
+      def channel_builder(channel_builder)
+        @channel_builder = channel_builder
+        self
+      end
+      alias channelBuilder channel_builder
+
+      # Accepted for PySpark compatibility; the master is chosen by the Spark
+      # Connect server, so this is a no-op.
+      # @return [self]
+      def master(_master)
+        self
+      end
+
+      # Hive support is a server-side setting and cannot be enabled from a
+      # Spark Connect client (PySpark raises too).
+      # @raise [NotImplementedError]
+      def enable_hive_support
+        raise NotImplementedError, "enableHiveSupport is not supported over Spark Connect; configure the server instead"
+      end
+      alias enableHiveSupport enable_hive_support
+
       # Set the application name.
       # @return [self]
       def app_name(name)
         @options["spark.app.name"] = name
         self
       end
+      alias appName app_name
 
       # Set an arbitrary configuration option to apply after connecting.
+      # @overload config(key, value)
+      # @overload config(map)
+      #   @param map [Hash] several options at once.
       # @return [self]
-      def config(key, value)
-        @options[key.to_s] = value
+      def config(key, value = nil)
+        if key.is_a?(Hash)
+          key.each { |k, v| @options[k.to_s] = v }
+        else
+          @options[key.to_s] = value
+        end
         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
+      # Reuse the active (or default) session, else build a new one. Builder
+      # options are applied to the returned session either way, like PySpark.
+      # With a block, yields the session and stops it afterwards (the Ruby
+      # counterpart of PySpark's `with ... as spark:`), returning the block's value.
+      #
+      # @yieldparam session [SparkSession]
+      # @return [SparkSession, Object]
+      def get_or_create(&block)
+        session = SparkSession.active
+        if session
+          apply_options(session)
+        else
+          session = create
+        end
+        block ? use(session, &block) : 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
+      # Always build a new {SparkSession} (which becomes the active session,
+      # and the default one when none is set). With a block, behaves like
+      # {#get_or_create} with a block.
+      #
+      # @yieldparam session [SparkSession]
+      # @return [SparkSession, Object]
+      def create(&block)
+        raise IllegalArgumentError, "Set either remote(...) or channel_builder(...), not both" if @remote && @channel_builder
+
+        channel = @channel_builder || ChannelBuilder.new(@remote || ENV["SPARK_REMOTE"] || "sc://localhost:15002")
+        session = SparkSession.new(SparkConnectClient.new(channel))
+        SparkSession.active = session
+        apply_options(session)
+        block ? use(session, &block) : session
       end
       alias build create
+
+      private
+
+      def apply_options(session)
+        opts = @options.except(*CONNECT_ONLY_OPTIONS)
+        session.conf.set_all(opts, silent: true) unless opts.empty?
+      end
+
+      def use(session)
+        yield session
+      ensure
+        session.stop
+      end
     end
   end
 end

data/lib/spark_connect/stat_functions.rb

@@ -22,10 +22,13 @@ module SparkConnect
       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"`).
+    # Correlation of two columns (`method` is `"pearson"`, the only supported one).
     # @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)
+    def corr(col1, col2, method = nil)
+      method = "pearson" if method.nil? || method.to_s.empty?
+      raise IllegalArgumentError, "corr method must be 'pearson', got #{method.inspect}" unless method.to_s == "pearson"
+
+      rel = Proto::StatCorr.new(input: @df.relation, col1: col1.to_s, col2: col2.to_s, method: method.to_s)
       scalar(@df.build(corr: rel))
     end
 
@@ -40,8 +43,9 @@ module SparkConnect
     # @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)
+    def freq_items(cols, support = nil)
+      support = 0.01 if support.nil? || support.zero?
+      rel = Proto::StatFreqItems.new(input: @df.relation, cols: Array(cols).map(&:to_s), support: support.to_f)
       @df.build(freq_items: rel)
     end
 
@@ -52,10 +56,15 @@ module SparkConnect
     # @param relative_error [Float]
     # @return [Array<Float>, Array<Array<Float>>] one list per column.
     def approx_quantile(cols, probabilities, relative_error)
+      unless probabilities.all? { |p| p.is_a?(Numeric) && p.between?(0, 1) }
+        raise IllegalArgumentError, "probabilities must be numbers in [0, 1], got #{probabilities.inspect}"
+      end
+      raise IllegalArgumentError, "relative_error must be non-negative, got #{relative_error}" if relative_error.negative?
+
       single = !cols.is_a?(Array)
       rel = Proto::StatApproxQuantile.new(
         input: @df.relation, cols: Array(cols).map(&:to_s),
-        probabilities: probabilities, relative_error: relative_error
+        probabilities: probabilities.map(&:to_f), relative_error: relative_error.to_f
       )
       row = @df.build(approx_quantile: rel).collect.first
       result = row.to_a
@@ -71,10 +80,11 @@ module SparkConnect
     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)
+        Proto::StatSampleBy::Fraction.new(stratum: Column.to_literal(stratum), fraction: fraction.to_f)
       end
-      rel = Proto::StatSampleBy.new(input: @df.relation, col: col_expr, fractions: frac)
-      rel.seed = seed if seed
+      # Like PySpark, pick the seed on the client so the sample is stable across actions.
+      rel = Proto::StatSampleBy.new(input: @df.relation, col: col_expr, fractions: frac,
+                                    seed: seed.nil? ? Random.rand(2**63) : Integer(seed))
       @df.build(sample_by: rel)
     end