Rename from raysql to datafusion_ray in multiple places (#13)

Author: timsaucer

README.md

@@ -46,15 +46,15 @@ import os
 import pandas as pd
 import ray
 
-from datafusion_ray import RaySqlContext
+from datafusion_ray import DatafusionRayContext
 
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
 
 # Start a local cluster
 ray.init(resources={"worker": 1})
 
 # Create a context and register a table
-ctx = RaySqlContext(2, use_ray_shuffle=True)
+ctx = DatafusionRayContext(2, use_ray_shuffle=True)
 # Register either a CSV or Parquet file
 # ctx.register_csv("tips", f"{SCRIPT_DIR}/tips.csv", True)
 ctx.register_parquet("tips", f"{SCRIPT_DIR}/tips.parquet")

build.rs

@@ -32,14 +32,14 @@ fn main() -> Result<(), String> {
 
     // We don't include the proto files in releases so that downstreams
     // do not need to have PROTOC included
-    if Path::new("src/proto/raysql.proto").exists() {
+    if Path::new("src/proto/datafusion-ray.proto").exists() {
         println!("cargo:rerun-if-changed=src/proto/datafusion.proto");
-        println!("cargo:rerun-if-changed=src/proto/raysql.proto");
+        println!("cargo:rerun-if-changed=src/proto/datafusion-ray.proto");
         tonic_build::configure()
             .extern_path(".datafusion", "::datafusion_proto::protobuf")
-            .compile(&["src/proto/raysql.proto"], &["src/proto"])
+            .compile(&["src/proto/datafusion-ray.proto"], &["src/proto"])
             .map_err(|e| format!("protobuf compilation failed: {e}"))?;
-        let generated_source_path = out.join("raysql.protobuf.rs");
+        let generated_source_path = out.join("datafusion-ray.protobuf.rs");
         let code = std::fs::read_to_string(generated_source_path).unwrap();
         let mut file = std::fs::OpenOptions::new()
             .write(true)

datafusion_ray/__init__.py

@@ -28,6 +28,6 @@
     serialize_execution_plan,
     deserialize_execution_plan,
 )
-from .context import RaySqlContext
+from .context import DatafusionRayContext
 
 __version__ = importlib_metadata.version(__name__)

datafusion_ray/context.py

@@ -27,6 +27,7 @@
 from datafusion_ray import Context, ExecutionGraph, QueryStage
 from typing import List
 
+
 def schedule_execution(
     graph: ExecutionGraph,
     stage_id: int,
@@ -208,7 +209,7 @@ def execute_query_partition(
     return ret[0] if len(ret) == 1 else ret
 
 
-class RaySqlContext:
+class DatafusionRayContext:
     def __init__(self, num_workers: int = 1, use_ray_shuffle: bool = False):
         self.ctx = Context(num_workers, use_ray_shuffle)
         self.num_workers = num_workers

datafusion_ray/main.py

@@ -20,7 +20,7 @@
 
 from pyarrow import csv as pacsv
 import ray
-from datafusion_ray import RaySqlContext
+from datafusion_ray import DatafusionRayContext
 
 NUM_CPUS_PER_WORKER = 8
 
@@ -31,9 +31,9 @@
 RESULTS_DIR = f"results-sf{SF}"
 
 
-def setup_context(use_ray_shuffle: bool, num_workers: int = 2) -> RaySqlContext:
+def setup_context(use_ray_shuffle: bool, num_workers: int = 2) -> DatafusionRayContext:
     print(f"Using {num_workers} workers")
-    ctx = RaySqlContext(num_workers, use_ray_shuffle)
+    ctx = DatafusionRayContext(num_workers, use_ray_shuffle)
     for table in [
         "customer",
         "lineitem",
@@ -53,14 +53,14 @@ def load_query(n: int) -> str:
         return fin.read()
 
 
-def tpch_query(ctx: RaySqlContext, q: int = 1):
+def tpch_query(ctx: DatafusionRayContext, q: int = 1):
     sql = load_query(q)
     result_set = ctx.sql(sql)
     return result_set
 
 
 def tpch_timing(
-    ctx: RaySqlContext,
+    ctx: DatafusionRayContext,
     q: int = 1,
     print_result: bool = False,
     write_result: bool = False,

examples/tips.py

@@ -19,15 +19,15 @@
 import pandas as pd
 import ray
 
-from raysql import RaySqlContext
+from datafusion_ray import DatafusionRayContext
 
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
 
 # Start a local cluster
 ray.init(resources={"worker": 1})
 
 # Create a context and register a table
-ctx = RaySqlContext(2, use_ray_shuffle=True)
+ctx = DatafusionRayContext(2, use_ray_shuffle=True)
 # Register either a CSV or Parquet file
 # ctx.register_csv("tips", f"{SCRIPT_DIR}/tips.csv", True)
 ctx.register_parquet("tips", f"{SCRIPT_DIR}/tips.parquet")

src/context.rs

@@ -44,7 +44,7 @@ use tokio::task::JoinHandle;
 
 type PyResultSet = Vec<PyObject>;
 
-#[pyclass(name = "Context", module = "raysql", subclass)]
+#[pyclass(name = "Context", module = "datafusion_ray", subclass)]
 pub struct PyContext {
     pub(crate) ctx: SessionContext,
     use_ray_shuffle: bool,

src/lib.rs

@@ -31,7 +31,7 @@ pub mod utils;
 
 /// A Python module implemented in Rust.
 #[pymodule]
-fn _datafusion_ray_internal(_py: Python, m: &PyModule) -> PyResult<()> {
+fn _datafusion_ray_internal(m: &Bound<'_, PyModule>) -> PyResult<()> {
     // register classes that can be created directly from Python code
     m.add_class::<context::PyContext>()?;
     m.add_class::<planner::PyExecutionGraph>()?;

src/planner.rs

@@ -32,7 +32,7 @@ use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use uuid::Uuid;
 
-#[pyclass(name = "ExecutionGraph", module = "raysql", subclass)]
+#[pyclass(name = "ExecutionGraph", module = "datafusion_ray", subclass)]
 pub struct PyExecutionGraph {
     pub graph: ExecutionGraph,
 }

src/proto/raysql.proto src/proto/datafusion_ray.proto

@@ -1,9 +1,9 @@
 syntax = "proto3";
 
-package raysql.protobuf;
+package datafusion_ray.protobuf;
 
 option java_multiple_files = true;
-option java_package = "raysql.protobuf";
+option java_package = "datafusion_ray.protobuf";
 option java_outer_classname = "RaySqlProto";
 
 import "datafusion.proto";