feat: add testing utilities for DataFrame range generation

Author: kosiew

python/tests/test_io.py

@@ -14,14 +14,16 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+
 from pathlib import Path
 
 import pyarrow as pa
 import pytest
 from datafusion import column
-from datafusion._testing import range_table
 from datafusion.io import read_avro, read_csv, read_json, read_parquet
 
+from .utils import range_table
+
 
 def test_read_json_global_ctx(ctx):
     path = Path(__file__).parent.resolve()

python/datafusion/_testing.py renamed to python/tests/utils.py

@@ -9,10 +9,9 @@
 
 from typing import TYPE_CHECKING
 
-from .context import SessionContext
-
 if TYPE_CHECKING:
     from datafusion import DataFrame
+    from datafusion.context import SessionContext
 
 
 def range_table(
@@ -42,5 +41,5 @@ def range_table(
         start, stop = 0, start
 
     parts = f", {int(partitions)}" if partitions is not None else ""
-    sql = f"SELECT * FROM range({int(start)}, {int(stop)}, {int(step)}{parts})"  # noqa: S608
+    sql = f"SELECT * FROM range({int(start)}, {int(stop)}, {int(step)}{parts})"
     return ctx.sql(sql)