[Chores] Update packagings

Signed-off-by: Austin Liu <[email protected]>

Author: austin362667

raysql/__init__.py datafusion_ray/__init__.py

@@ -20,7 +20,7 @@
 except ImportError:
     import importlib_metadata
 
-from ._raysql_internal import (
+from ._datafusion_ray_internal import (
     Context,
     ExecutionGraph,
     QueryStage,

raysql/context.py datafusion_ray/context.py

@@ -23,8 +23,8 @@
 import pyarrow as pa
 import ray
 
-import raysql
-from raysql import Context, ExecutionGraph, QueryStage
+import datafusion_ray
+from datafusion_ray import Context, ExecutionGraph, QueryStage
 from typing import List
 
 def schedule_execution(
@@ -73,7 +73,7 @@ def _get_worker_inputs(
         return ids, futures
 
     # schedule the actual execution workers
-    plan_bytes = raysql.serialize_execution_plan(stage.get_execution_plan())
+    plan_bytes = datafusion_ray.serialize_execution_plan(stage.get_execution_plan())
     futures = []
     opt = {}
     opt["resources"] = {"worker": 1e-3}
@@ -153,7 +153,7 @@ def _get_worker_inputs(
         ray.get([f for _, lst in child_outputs for f in lst])
 
     # schedule the actual execution workers
-    plan_bytes = raysql.serialize_execution_plan(stage.get_execution_plan())
+    plan_bytes = datafusion_ray.serialize_execution_plan(stage.get_execution_plan())
     futures = []
     opt = {}
     opt["resources"] = {"worker": 1e-3}
@@ -179,7 +179,7 @@ def execute_query_partition(
     *input_partitions: list[pa.RecordBatch],
 ) -> Iterable[pa.RecordBatch]:
     start_time = time.time()
-    plan = raysql.deserialize_execution_plan(plan_bytes)
+    plan = datafusion_ray.deserialize_execution_plan(plan_bytes)
     # print(
     #     "Worker executing plan {} partition #{} with shuffle inputs {}".format(
     #         plan.display(),
@@ -193,7 +193,7 @@ def execute_query_partition(
     # This is delegating to DataFusion for execution, but this would be a good place
     # to plug in other execution engines by translating the plan into another engine's plan
     # (perhaps via Substrait, once DataFusion supports converting a physical plan to Substrait)
-    ret = raysql.execute_partition(plan, part, partitions)
+    ret = datafusion_ray.execute_partition(plan, part, partitions)
     duration = time.time() - start_time
     event = {
         "cat": f"{stage_id}-{part}",
@@ -238,7 +238,7 @@ def sql(self, sql: str) -> pa.RecordBatch:
         else:
             # serialize the query stages and store in Ray object store
             query_stages = [
-                raysql.serialize_execution_plan(
+                datafusion_ray.serialize_execution_plan(
                     graph.get_query_stage(i).get_execution_plan()
                 )
                 for i in range(final_stage_id + 1)

raysql/main.py datafusion_ray/main.py

@@ -20,18 +20,15 @@
 
 from pyarrow import csv as pacsv
 import ray
-from raysql import RaySqlContext
+from datafusion_ray import RaySqlContext
 
 NUM_CPUS_PER_WORKER = 8
 
-SF = 10
+SF = 1
 DATA_DIR = f"/mnt/data0/tpch/sf{SF}-parquet"
 SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
 QUERIES_DIR = os.path.join(SCRIPT_DIR, f"../sqlbench-h/queries/sf={SF}")
 RESULTS_DIR = f"results-sf{SF}"
-TRUTH_DIR = (
-    "/home/ubuntu/raysort/ray-sql/sqlbench-runners/spark/{RESULTS_DIR}/{RESULTS_DIR}"
-)
 
 
 def setup_context(use_ray_shuffle: bool, num_workers: int = 2) -> RaySqlContext:
@@ -104,7 +101,7 @@ def compare(q: int):
 
 
 def tpch_bench():
-    ray.init("auto")
+    ray.init(resources={"worker": 1})
     num_workers = int(ray.cluster_resources().get("worker", 1)) * NUM_CPUS_PER_WORKER
     use_ray_shuffle = False
     ctx = setup_context(use_ray_shuffle, num_workers)

raysql/ray_utils.py datafusion_ray/ray_utils.py

@@ -16,7 +16,7 @@
 # under the License.
 
 import pytest
-from raysql import Context
+from datafusion_ray import Context
 
 def test():
     ctx = Context(1, False)