added a BallistaContext to ballista to allow for Remote or standalone (…

…#1100) * added a pycontext to ballista * added a pycontext to ballista * added a pycontext to ballista * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updating the pyballista package to ballista * changing the packagaing naming convention from pyballista to ballista * changing the packagaing naming convention from pyballista to ballista * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * updated python to have two static methods for creating a ballista context * Updating BallistaContext and Config * Updating BallistaContext and Config * updated python to have two static methods for creating a ballista context * Updating BallistaContext and Config, calling it for the night, will complete tomorrow * Updating BallistaContext and Config, calling it for the night, will complete tomorrow * Adding config to ballista context * Adding config to ballista context * Adding config to ballista context * Adding config to ballista context * Updated Builder and Docs * Updated Builder and Docs * Updated Builder and Docs * Updated Builder and Docs * Updated Builder and Docs * Updated Builder and Docs --------- Co-authored-by: Trevor Barnes <[email protected]>
apache · Nov 19, 2024 · d949e5f · d949e5f
1 parent 5b6b50b
commit d949e5f
Show file tree

Hide file tree

Showing 11 changed files with 150 additions and 389 deletions.
diff --git a/docs/source/user-guide/python.md b/docs/source/user-guide/python.md
@@ -28,9 +28,20 @@ popular file formats files, run it in a distributed environment, and obtain the
 
 The following code demonstrates how to create a Ballista context and connect to a scheduler.
 
+If you are running a standalone cluster (runs locally), all you need to do is call the stand alone cluster method `standalone()` or your BallistaContext. If you are running a cluster in remote mode, you need to provide the URL `Ballista.remote("http://my-remote-ip:50050")`.
+
 ```text
->>> import ballista
->>> ctx = ballista.BallistaContext("localhost", 50050)
+>>> from ballista import BallistaBuilder
+>>> # for a standalone instance
+>>> # Ballista will initiate with an empty config
+>>> # set config variables with `config()`
+>>> ballista = BallistaBuilder()\
+>>>    .config("ballista.job.name", "example ballista")
+>>>
+>>> ctx = ballista.standalone()
+>>>
+>>> # for a remote instance provide the URL
+>>> ctx = ballista.remote("df://url-path-to-scheduler:50050")
 ```
 
 ## SQL
@@ -103,14 +114,15 @@ The `explain` method can be used to show the logical and physical query plans fo
 The following example demonstrates creating arrays with PyArrow and then creating a Ballista DataFrame.
 
 ```python
-import ballista
+from ballista import BallistaBuilder
 import pyarrow
 
 # an alias
+# TODO implement Functions
 f = ballista.functions
 
 # create a context
-ctx = ballista.BallistaContext("localhost", 50050)
+ctx = Ballista().standalone()
 
 # create a RecordBatch and a new DataFrame from it
 batch = pyarrow.RecordBatch.from_arrays(

diff --git a/python/Cargo.toml b/python/Cargo.toml
@@ -26,14 +26,14 @@ readme = "README.md"
 license = "Apache-2.0"
 edition = "2021"
 rust-version = "1.72"
-include = ["/src", "/pyballista", "/LICENSE.txt", "pyproject.toml", "Cargo.toml", "Cargo.lock"]
+include = ["/src", "/ballista", "/LICENSE.txt", "pyproject.toml", "Cargo.toml", "Cargo.lock"]
 publish = false
 
 [dependencies]
 async-trait = "0.1.77"
-ballista = { path = "../ballista/client", version = "0.12.0" }
+ballista = { path = "../ballista/client", version = "0.12.0", features = ["standalone"] }
 ballista-core = { path = "../ballista/core", version = "0.12.0" }
-datafusion = { version = "42" }
+datafusion = { version = "42", features = ["pyarrow", "avro"] }
 datafusion-proto = { version = "42" }
 datafusion-python = { version = "42" }
 
@@ -43,6 +43,4 @@ tokio = { version = "1.35", features = ["macros", "rt", "rt-multi-thread", "sync
 
 [lib]
 crate-type = ["cdylib"]
-name = "pyballista"
-
-
+name = "ballista"
diff --git a/python/README.md b/python/README.md
@@ -29,8 +29,8 @@ part of the default Cargo workspace so that it doesn't cause overhead for mainta
 Creates a new context and connects to a Ballista scheduler process.
 
 ```python
-from pyballista import SessionContext
->>> ctx = SessionContext("localhost", 50050)
+from ballista import BallistaBuilder
+>>> ctx = BallistaBuilder().standalone()
 ```
 
 ## Example SQL Usage

diff --git a/python/pyballista/__init__.py → python/ballista/__init__.py b/python/pyballista/__init__.py → python/ballista/__init__.py
@@ -25,12 +25,12 @@
 
 import pyarrow as pa
 
-from .pyballista_internal import (
-    SessionContext,
+from .ballista_internal import (
+    BallistaBuilder,
 )
 
 __version__ = importlib_metadata.version(__name__)
 
 __all__ = [
-    "SessionContext",
-]
+    "BallistaBuilder",
+]
diff --git a/python/pyballista/tests/__init__.py → python/ballista/tests/__init__.py b/python/pyballista/tests/__init__.py → python/ballista/tests/__init__.py
diff --git a/python/pyballista/tests/test_context.py → python/ballista/tests/test_context.py b/python/pyballista/tests/test_context.py → python/ballista/tests/test_context.py
@@ -15,50 +15,50 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from pyballista import SessionContext
+from ballista import BallistaBuilder
 import pytest
 
 def test_create_context():
-    SessionContext("localhost", 50050)
+    BallistaBuilder().standalone()
 
 def test_select_one():
-    ctx = SessionContext("localhost", 50050)
+    ctx = BallistaBuilder().standalone()
     df = ctx.sql("SELECT 1")
     batches = df.collect()
     assert len(batches) == 1
 
 def test_read_csv():
-    ctx = SessionContext("localhost", 50050)
+    ctx = BallistaBuilder().standalone()
     df = ctx.read_csv("testdata/test.csv", has_header=True)
     batches = df.collect()
     assert len(batches) == 1
     assert len(batches[0]) == 1
 
 def test_register_csv():
-    ctx = SessionContext("localhost", 50050)
+    ctx = BallistaBuilder().standalone()
     ctx.register_csv("test", "testdata/test.csv", has_header=True)
     df = ctx.sql("SELECT * FROM test")
     batches = df.collect()
     assert len(batches) == 1
     assert len(batches[0]) == 1
 
 def test_read_parquet():
-    ctx = SessionContext("localhost", 50050)
+    ctx = BallistaBuilder().standalone()
     df = ctx.read_parquet("testdata/test.parquet")
     batches = df.collect()
     assert len(batches) == 1
     assert len(batches[0]) == 8
 
 def test_register_parquet():
-    ctx = SessionContext("localhost", 50050)
+    ctx = BallistaBuilder().standalone()
     ctx.register_parquet("test", "testdata/test.parquet")
     df = ctx.sql("SELECT * FROM test")
     batches = df.collect()
     assert len(batches) == 1
     assert len(batches[0]) == 8
 
 def test_read_dataframe_api():
-    ctx = SessionContext("localhost", 50050)
+    ctx = BallistaBuilder().standalone()
     df = ctx.read_csv("testdata/test.csv", has_header=True) \
         .select_columns('a', 'b') \
         .limit(1)
@@ -67,11 +67,12 @@ def test_read_dataframe_api():
     assert len(batches[0]) == 1
 
 def test_execute_plan():
-    ctx = SessionContext("localhost", 50050)
+    ctx = BallistaBuilder().standalone()
     df = ctx.read_csv("testdata/test.csv", has_header=True) \
         .select_columns('a', 'b') \
         .limit(1)
-    df = ctx.execute_logical_plan(df.logical_plan())
+    # TODO research SessionContext Logical Plan for DataFusionPython
+    #df = ctx.execute_logical_plan(df.logical_plan())
     batches = df.collect()
     assert len(batches) == 1
     assert len(batches[0]) == 1
diff --git a/python/examples/example.py b/python/examples/example.py
@@ -0,0 +1,32 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from ballista import BallistaBuilder
+from datafusion.context import SessionContext
+
+# Ballista will initiate with an empty config
+# set config variables with `config`
+ctx: SessionContext = BallistaBuilder()\
+    .config("ballista.job.name", "example ballista")\
+    .config("ballista.shuffle.partitions", "16")\
+    .standalone()
+
+#ctx_remote: SessionContext = ballista.remote("remote_ip", 50050)
+
+# Select 1 to verify its working
+ctx.sql("SELECT 1").show()
+#ctx_remote.sql("SELECT 2").show()
diff --git a/python/pyproject.toml b/python/pyproject.toml
@@ -20,7 +20,7 @@ requires = ["maturin>=0.15,<0.16"]
 build-backend = "maturin"
 
 [project]
-name = "pyballista"
+name = "ballista"
 description = "Python client for Apache Arrow Ballista Distributed SQL Query Engine"
 readme = "README.md"
 license = {file = "LICENSE.txt"}
@@ -55,10 +55,10 @@ repository = "https://github.com/apache/arrow-ballista"
 profile = "black"
 
 [tool.maturin]
-module-name = "pyballista.pyballista_internal"
+module-name = "ballista.ballista_internal"
 include = [
     { path = "Cargo.lock", format = "sdist" }
 ]
 exclude = [".github/**", "ci/**", ".asf.yaml"]
 # Require Cargo.lock is up to date
-locked = true
+locked = true