andygrove commented on code in PR #2748:
URL:
https://github.com/apache/datafusion-ballista/pull/2748#discussion_r4460635590
##########
python/src/cluster.rs:
##########
@@ -53,7 +58,10 @@ impl PyScheduler {
PyBallistaPhysicalExtensionCodec::default(),
)),
..Default::default()
- };
+ }
+ // S3 support, as in the scheduler binary
+
.with_override_config_producer(Arc::new(session_config_with_s3_support))
+
.with_override_session_builder(Arc::new(session_state_with_s3_support));
Review Comment:
Switching to `session_state_with_s3_support` means the scheduler session no
longer gets `ballista_default_optimizer_rules()`, which
`default_session_builder` installs (`ballista/core/src/utils.rs:221`). That
drops `NotInSubqueryRewrite`, so a `WHERE x NOT IN (SELECT ...)` query would be
planned by DataFusion as a null-aware anti join, which is the shape the rule is
there to avoid.
I think the Rust scheduler binary has the same gap today, so this PR spreads
it to the Python scheduler and the Python test cluster (line 311 has the same
change). Would it make sense to add the optimizer rules inside
`session_state_with_s3_support` so every caller gets them? A quick `NOT IN`
test on `setup_test_cluster()` would confirm either way.
##########
python/src/cluster.rs:
##########
@@ -191,6 +199,14 @@ impl PyExecutor {
config.vcores = vcores as usize;
}
+ // S3 support, as in the executor binary
+ config.override_config_producer =
Some(Arc::new(session_config_with_s3_support));
+ let s3_runtime: RuntimeProducer =
Arc::new(runtime_env_with_s3_support);
+ config.override_runtime_producer = Some(match config.work_dir.clone() {
+ Some(work_dir) => with_spill_dir(s3_runtime, work_dir),
Review Comment:
`PyExecutor::new` never sets `config.work_dir`, so it is always `None` here
and this arm can't run. The effect is that Python executors now spill to the OS
temp dir, where the default runtime used to spill under the work dir that gets
cleaned up with the executor. The comment saying this matches the executor
binary is also a bit misleading for that reason.
Either exposing `work_dir` on the constructor or dropping the branch would
fix it. Another option is for `start_executor_process` to apply the spill dir
itself once it has resolved the work dir, which would help every caller that
sets an override runtime producer.
##########
python/python/tests/test_s3.py:
##########
@@ -0,0 +1,138 @@
+# 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.
+
+"""
+S3 tests against a RustFS container, skipped when docker is not available.
+
+RustFS replaces MinIO, which withdrew its images, same as the Rust
+integration tests in `examples/tests/common/mod.rs`.
+"""
+
+from ballista import BallistaSessionContext, setup_test_cluster
+from datafusion import col, udf
+from datafusion.object_store import AmazonS3
+import pyarrow as pa
+import pytest
+
+pytest.importorskip("testcontainers")
+
+from testcontainers.core.container import DockerContainer # noqa: E402
+from testcontainers.core.wait_strategies import HttpWaitStrategy # noqa: E402
+
+TEST_PARQUET = "testdata/test.parquet"
Review Comment:
This path is relative, so running pytest from anywhere other than `python/`
makes the fixture error out instead of skipping. Anchoring it to
`Path(__file__)` would avoid that.
##########
python/python/tests/test_s3.py:
##########
@@ -0,0 +1,138 @@
+# 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.
+
+"""
+S3 tests against a RustFS container, skipped when docker is not available.
+
+RustFS replaces MinIO, which withdrew its images, same as the Rust
+integration tests in `examples/tests/common/mod.rs`.
+"""
+
+from ballista import BallistaSessionContext, setup_test_cluster
+from datafusion import col, udf
+from datafusion.object_store import AmazonS3
+import pyarrow as pa
+import pytest
+
+pytest.importorskip("testcontainers")
+
+from testcontainers.core.container import DockerContainer # noqa: E402
+from testcontainers.core.wait_strategies import HttpWaitStrategy # noqa: E402
+
+TEST_PARQUET = "testdata/test.parquet"
+
+RUSTFS_IMAGE = "rustfs/rustfs:1.0.0"
+S3_PORT = 9000
+REGION = "eu-west-1"
+BUCKET = "ballista"
+ACCESS_KEY_ID = "BALLISTA"
+SECRET_KEY = "BALLISTASECRET"
+
+
+def _docker_available() -> bool:
+ try:
+ import docker
+
+ return docker.from_env().ping()
+ except Exception:
+ return False
+
+
+pytestmark = pytest.mark.skipif(
+ not _docker_available(), reason="docker is required for S3 tests"
+)
+
+
[email protected](scope="module")
+def s3_endpoint():
+ # RustFS logs only warnings by default, wait on its health endpoint
+ # rather than a startup message
+ container = (
+ DockerContainer(RUSTFS_IMAGE)
+ .with_exposed_ports(S3_PORT)
+ .with_env("RUSTFS_ACCESS_KEY", ACCESS_KEY_ID)
+ .with_env("RUSTFS_SECRET_KEY", SECRET_KEY)
+ .waiting_for(
+ HttpWaitStrategy(S3_PORT, "/health")
+ .for_status_code(200)
+ .with_startup_timeout(60)
+ )
+ )
+ with container:
+ # a top-level directory under the data volume is a bucket
+ exit_code, output = container.exec(["mkdir", f"/data/{BUCKET}"])
+ assert exit_code == 0, output
+
+ host = container.get_container_host_ip()
+ port = container.get_exposed_port(S3_PORT)
+ yield f"http://{host}:{port}";
+
+
[email protected](scope="module")
+def ctx(s3_endpoint):
+ host, port = setup_test_cluster()
+ ctx = BallistaSessionContext(f"df://{host}:{port}")
+
+ # used by the client for schema inference
+ ctx.register_object_store(
+ "s3://",
+ AmazonS3(
+ bucket_name=BUCKET,
+ region=REGION,
+ access_key_id=ACCESS_KEY_ID,
+ secret_access_key=SECRET_KEY,
+ endpoint=s3_endpoint,
+ allow_http=True,
+ ),
+ BUCKET,
+ )
+
+ # used by the scheduler and executors
+ ctx.sql(f"SET s3.region = '{REGION}'")
Review Comment:
A question on the client side of the S3 story. The Python client doesn't
seem to register `S3Options` itself, so could you confirm that `SET s3.*` works
on a plain `BallistaSessionContext`? It is probably fine since the Python CI
jobs are green. If so, a short note in `quickstart.md` on how to supply S3
credentials would help users, because right now the docs don't say.
##########
python/python/ballista/jupyter.py:
##########
@@ -534,14 +533,17 @@ def _show_help(self) -> Optional[str]:
def explain_visual(df: DataFrame, analyze: bool = False) ->
ExecutionPlanVisualization:
Review Comment:
The docstring now says `analyze` only labels the visualization. Since it
does nothing else, would it be worth deprecating the parameter, or at least
warning when it is `True`? Someone calling `explain_visual(df, analyze=True)`
will expect runtime statistics and quietly get the static logical plan instead.
##########
python/python/tests/test_s3.py:
##########
@@ -0,0 +1,138 @@
+# 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.
+
+"""
+S3 tests against a RustFS container, skipped when docker is not available.
+
+RustFS replaces MinIO, which withdrew its images, same as the Rust
+integration tests in `examples/tests/common/mod.rs`.
+"""
+
+from ballista import BallistaSessionContext, setup_test_cluster
+from datafusion import col, udf
+from datafusion.object_store import AmazonS3
+import pyarrow as pa
+import pytest
+
+pytest.importorskip("testcontainers")
+
+from testcontainers.core.container import DockerContainer # noqa: E402
+from testcontainers.core.wait_strategies import HttpWaitStrategy # noqa: E402
+
+TEST_PARQUET = "testdata/test.parquet"
+
+RUSTFS_IMAGE = "rustfs/rustfs:1.0.0"
+S3_PORT = 9000
+REGION = "eu-west-1"
+BUCKET = "ballista"
+ACCESS_KEY_ID = "BALLISTA"
+SECRET_KEY = "BALLISTASECRET"
+
+
+def _docker_available() -> bool:
+ try:
+ import docker
+
+ return docker.from_env().ping()
+ except Exception:
+ return False
+
+
+pytestmark = pytest.mark.skipif(
+ not _docker_available(), reason="docker is required for S3 tests"
+)
+
+
[email protected](scope="module")
+def s3_endpoint():
+ # RustFS logs only warnings by default, wait on its health endpoint
+ # rather than a startup message
+ container = (
+ DockerContainer(RUSTFS_IMAGE)
+ .with_exposed_ports(S3_PORT)
+ .with_env("RUSTFS_ACCESS_KEY", ACCESS_KEY_ID)
+ .with_env("RUSTFS_SECRET_KEY", SECRET_KEY)
+ .waiting_for(
+ HttpWaitStrategy(S3_PORT, "/health")
+ .for_status_code(200)
+ .with_startup_timeout(60)
+ )
+ )
+ with container:
+ # a top-level directory under the data volume is a bucket
+ exit_code, output = container.exec(["mkdir", f"/data/{BUCKET}"])
+ assert exit_code == 0, output
+
+ host = container.get_container_host_ip()
+ port = container.get_exposed_port(S3_PORT)
+ yield f"http://{host}:{port}";
+
+
[email protected](scope="module")
+def ctx(s3_endpoint):
+ host, port = setup_test_cluster()
+ ctx = BallistaSessionContext(f"df://{host}:{port}")
+
+ # used by the client for schema inference
+ ctx.register_object_store(
+ "s3://",
+ AmazonS3(
+ bucket_name=BUCKET,
+ region=REGION,
+ access_key_id=ACCESS_KEY_ID,
+ secret_access_key=SECRET_KEY,
+ endpoint=s3_endpoint,
+ allow_http=True,
+ ),
+ BUCKET,
+ )
+
+ # used by the scheduler and executors
+ ctx.sql(f"SET s3.region = '{REGION}'")
+ ctx.sql(f"SET s3.access_key_id = '{ACCESS_KEY_ID}'")
+ ctx.sql(f"SET s3.secret_access_key = '{SECRET_KEY}'")
+ ctx.sql(f"SET s3.endpoint = '{s3_endpoint}'")
+ ctx.sql("SET s3.allow_http = true")
+
+ # executors write the data, so the cluster has to have S3 support
+ ctx.read_parquet(TEST_PARQUET).write_parquet(f"s3://{BUCKET}/test/")
+ ctx.sql(f"CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION
's3://{BUCKET}/test/'")
+ return ctx
+
+
+def test_read_s3(ctx):
+ expected = ctx.read_parquet(TEST_PARQUET).sort(col("id")).to_arrow_table()
+
+ result = ctx.sql("SELECT * FROM t ORDER BY id").to_arrow_table()
+
+ assert result == expected
Review Comment:
Small thing: this compares tables read through different paths (the local
file versus the S3-written copy). If schema metadata or string types ever
differ it could flake. Comparing `.to_pylist()` or using `Table.equals(...,
check_metadata=False)` would be more forgiving.
##########
python/Cargo.toml:
##########
@@ -34,33 +34,27 @@ publish = false
async-ffi = { version = "0.5" }
async-trait = { version = "0.1.92" }
ballista = { path = "../ballista/client", version = "55.0.0" }
-ballista-core = { path = "../ballista/core", version = "55.0.0" }
+ballista-core = { path = "../ballista/core", version = "55.0.0", features =
["build-binary"] }
Review Comment:
Non-blocking. `build-binary` pulls `clap`, `lru` and the AWS and HTTP
`object_store` features into the wheel just to reach
`ballista_core::object_store`. That is fine for now. As a follow-up it might be
worth a narrower `object-store` feature in `ballista-core` that `build-binary`
also enables, so changes to the binary feature set don't silently change the
wheel.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
