Enabled ruff rule ICN001

Author: Spaarsh

benchmarks/db-benchmark/groupby-datafusion.py

@@ -20,7 +20,7 @@
 import timeit
 
 import datafusion as df
-import pyarrow
+import pyarrow as pa
 from datafusion import (
     RuntimeEnvBuilder,
     SessionConfig,
@@ -68,14 +68,14 @@ def execute(df) -> list:
 src_grp = os.path.join("data", data_name + ".csv")
 print("loading dataset %s" % src_grp, flush=True)
 
-schema = pyarrow.schema(
+schema = pa.schema(
     [
-        ("id4", pyarrow.int32()),
-        ("id5", pyarrow.int32()),
-        ("id6", pyarrow.int32()),
-        ("v1", pyarrow.int32()),
-        ("v2", pyarrow.int32()),
-        ("v3", pyarrow.float64()),
+        ("id4", pa.int32()),
+        ("id5", pa.int32()),
+        ("id6", pa.int32()),
+        ("v1", pa.int32()),
+        ("v2", pa.int32()),
+        ("v3", pa.float64()),
     ]
 )
 

examples/python-udaf.py

@@ -16,7 +16,7 @@
 # under the License.
 
 import datafusion
-import pyarrow
+import pyarrow as pa
 import pyarrow.compute
 from datafusion import Accumulator, col, udaf
 
@@ -27,47 +27,43 @@ class MyAccumulator(Accumulator):
     """
 
     def __init__(self) -> None:
-        self._sum = pyarrow.scalar(0.0)
+        self._sum = pa.scalar(0.0)
 
-    def update(self, values: pyarrow.Array) -> None:
+    def update(self, values: pa.Array) -> None:
         # not nice since pyarrow scalars can't be summed yet. This breaks on `None`
-        self._sum = pyarrow.scalar(
-            self._sum.as_py() + pyarrow.compute.sum(values).as_py()
-        )
+        self._sum = pa.scalar(self._sum.as_py() + pa.compute.sum(values).as_py())
 
-    def merge(self, states: pyarrow.Array) -> None:
+    def merge(self, states: pa.Array) -> None:
         # not nice since pyarrow scalars can't be summed yet. This breaks on `None`
-        self._sum = pyarrow.scalar(
-            self._sum.as_py() + pyarrow.compute.sum(states).as_py()
-        )
+        self._sum = pa.scalar(self._sum.as_py() + pa.compute.sum(states).as_py())
 
-    def state(self) -> pyarrow.Array:
-        return pyarrow.array([self._sum.as_py()])
+    def state(self) -> pa.Array:
+        return pa.array([self._sum.as_py()])
 
-    def evaluate(self) -> pyarrow.Scalar:
+    def evaluate(self) -> pa.Scalar:
         return self._sum
 
 
 # create a context
 ctx = datafusion.SessionContext()
 
 # create a RecordBatch and a new DataFrame from it
-batch = pyarrow.RecordBatch.from_arrays(
-    [pyarrow.array([1, 2, 3]), pyarrow.array([4, 5, 6])],
+batch = pa.RecordBatch.from_arrays(
+    [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
     names=["a", "b"],
 )
 df = ctx.create_dataframe([[batch]])
 
 my_udaf = udaf(
     MyAccumulator,
-    pyarrow.float64(),
-    pyarrow.float64(),
-    [pyarrow.float64()],
+    pa.float64(),
+    pa.float64(),
+    [pa.float64()],
     "stable",
 )
 
 df = df.aggregate([], [my_udaf(col("a"))])
 
 result = df.collect()[0]
 
-assert result.column(0) == pyarrow.array([6.0])
+assert result.column(0) == pa.array([6.0])

examples/python-udf.py

@@ -15,23 +15,23 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import pyarrow
+import pyarrow as pa
 from datafusion import SessionContext, udf
 from datafusion import functions as f
 
 
-def is_null(array: pyarrow.Array) -> pyarrow.Array:
+def is_null(array: pa.Array) -> pa.Array:
     return array.is_null()
 
 
-is_null_arr = udf(is_null, [pyarrow.int64()], pyarrow.bool_(), "stable")
+is_null_arr = udf(is_null, [pa.int64()], pa.bool_(), "stable")
 
 # create a context
 ctx = SessionContext()
 
 # create a RecordBatch and a new DataFrame from it
-batch = pyarrow.RecordBatch.from_arrays(
-    [pyarrow.array([1, 2, 3]), pyarrow.array([4, 5, 6])],
+batch = pa.RecordBatch.from_arrays(
+    [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
     names=["a", "b"],
 )
 df = ctx.create_dataframe([[batch]])
@@ -40,4 +40,4 @@ def is_null(array: pyarrow.Array) -> pyarrow.Array:
 
 result = df.collect()[0]
 
-assert result.column(0) == pyarrow.array([False] * 3)
+assert result.column(0) == pa.array([False] * 3)

examples/query-pyarrow-data.py

@@ -16,15 +16,15 @@
 # under the License.
 
 import datafusion
-import pyarrow
+import pyarrow as pa
 from datafusion import col
 
 # create a context
 ctx = datafusion.SessionContext()
 
 # create a RecordBatch and a new DataFrame from it
-batch = pyarrow.RecordBatch.from_arrays(
-    [pyarrow.array([1, 2, 3]), pyarrow.array([4, 5, 6])],
+batch = pa.RecordBatch.from_arrays(
+    [pa.array([1, 2, 3]), pa.array([4, 5, 6])],
     names=["a", "b"],
 )
 df = ctx.create_dataframe([[batch]])
@@ -38,5 +38,5 @@
 # execute and collect the first (and only) batch
 result = df.collect()[0]
 
-assert result.column(0) == pyarrow.array([5, 7, 9])
-assert result.column(1) == pyarrow.array([-3, -3, -3])
+assert result.column(0) == pa.array([5, 7, 9])
+assert result.column(1) == pa.array([-3, -3, -3])

examples/tpch/convert_data_to_parquet.py

@@ -25,95 +25,95 @@
 import os
 
 import datafusion
-import pyarrow
+import pyarrow as pa
 
 ctx = datafusion.SessionContext()
 
 all_schemas = {}
 
 all_schemas["customer"] = [
-    ("C_CUSTKEY", pyarrow.int64()),
-    ("C_NAME", pyarrow.string()),
-    ("C_ADDRESS", pyarrow.string()),
-    ("C_NATIONKEY", pyarrow.int64()),
-    ("C_PHONE", pyarrow.string()),
-    ("C_ACCTBAL", pyarrow.decimal128(15, 2)),
-    ("C_MKTSEGMENT", pyarrow.string()),
-    ("C_COMMENT", pyarrow.string()),
+    ("C_CUSTKEY", pa.int64()),
+    ("C_NAME", pa.string()),
+    ("C_ADDRESS", pa.string()),
+    ("C_NATIONKEY", pa.int64()),
+    ("C_PHONE", pa.string()),
+    ("C_ACCTBAL", pa.decimal128(15, 2)),
+    ("C_MKTSEGMENT", pa.string()),
+    ("C_COMMENT", pa.string()),
 ]
 
 all_schemas["lineitem"] = [
-    ("L_ORDERKEY", pyarrow.int64()),
-    ("L_PARTKEY", pyarrow.int64()),
-    ("L_SUPPKEY", pyarrow.int64()),
-    ("L_LINENUMBER", pyarrow.int32()),
-    ("L_QUANTITY", pyarrow.decimal128(15, 2)),
-    ("L_EXTENDEDPRICE", pyarrow.decimal128(15, 2)),
-    ("L_DISCOUNT", pyarrow.decimal128(15, 2)),
-    ("L_TAX", pyarrow.decimal128(15, 2)),
-    ("L_RETURNFLAG", pyarrow.string()),
-    ("L_LINESTATUS", pyarrow.string()),
-    ("L_SHIPDATE", pyarrow.date32()),
-    ("L_COMMITDATE", pyarrow.date32()),
-    ("L_RECEIPTDATE", pyarrow.date32()),
-    ("L_SHIPINSTRUCT", pyarrow.string()),
-    ("L_SHIPMODE", pyarrow.string()),
-    ("L_COMMENT", pyarrow.string()),
+    ("L_ORDERKEY", pa.int64()),
+    ("L_PARTKEY", pa.int64()),
+    ("L_SUPPKEY", pa.int64()),
+    ("L_LINENUMBER", pa.int32()),
+    ("L_QUANTITY", pa.decimal128(15, 2)),
+    ("L_EXTENDEDPRICE", pa.decimal128(15, 2)),
+    ("L_DISCOUNT", pa.decimal128(15, 2)),
+    ("L_TAX", pa.decimal128(15, 2)),
+    ("L_RETURNFLAG", pa.string()),
+    ("L_LINESTATUS", pa.string()),
+    ("L_SHIPDATE", pa.date32()),
+    ("L_COMMITDATE", pa.date32()),
+    ("L_RECEIPTDATE", pa.date32()),
+    ("L_SHIPINSTRUCT", pa.string()),
+    ("L_SHIPMODE", pa.string()),
+    ("L_COMMENT", pa.string()),
 ]
 
 all_schemas["nation"] = [
-    ("N_NATIONKEY", pyarrow.int64()),
-    ("N_NAME", pyarrow.string()),
-    ("N_REGIONKEY", pyarrow.int64()),
-    ("N_COMMENT", pyarrow.string()),
+    ("N_NATIONKEY", pa.int64()),
+    ("N_NAME", pa.string()),
+    ("N_REGIONKEY", pa.int64()),
+    ("N_COMMENT", pa.string()),
 ]
 
 all_schemas["orders"] = [
-    ("O_ORDERKEY", pyarrow.int64()),
-    ("O_CUSTKEY", pyarrow.int64()),
-    ("O_ORDERSTATUS", pyarrow.string()),
-    ("O_TOTALPRICE", pyarrow.decimal128(15, 2)),
-    ("O_ORDERDATE", pyarrow.date32()),
-    ("O_ORDERPRIORITY", pyarrow.string()),
-    ("O_CLERK", pyarrow.string()),
-    ("O_SHIPPRIORITY", pyarrow.int32()),
-    ("O_COMMENT", pyarrow.string()),
+    ("O_ORDERKEY", pa.int64()),
+    ("O_CUSTKEY", pa.int64()),
+    ("O_ORDERSTATUS", pa.string()),
+    ("O_TOTALPRICE", pa.decimal128(15, 2)),
+    ("O_ORDERDATE", pa.date32()),
+    ("O_ORDERPRIORITY", pa.string()),
+    ("O_CLERK", pa.string()),
+    ("O_SHIPPRIORITY", pa.int32()),
+    ("O_COMMENT", pa.string()),
 ]
 
 all_schemas["part"] = [
-    ("P_PARTKEY", pyarrow.int64()),
-    ("P_NAME", pyarrow.string()),
-    ("P_MFGR", pyarrow.string()),
-    ("P_BRAND", pyarrow.string()),
-    ("P_TYPE", pyarrow.string()),
-    ("P_SIZE", pyarrow.int32()),
-    ("P_CONTAINER", pyarrow.string()),
-    ("P_RETAILPRICE", pyarrow.decimal128(15, 2)),
-    ("P_COMMENT", pyarrow.string()),
+    ("P_PARTKEY", pa.int64()),
+    ("P_NAME", pa.string()),
+    ("P_MFGR", pa.string()),
+    ("P_BRAND", pa.string()),
+    ("P_TYPE", pa.string()),
+    ("P_SIZE", pa.int32()),
+    ("P_CONTAINER", pa.string()),
+    ("P_RETAILPRICE", pa.decimal128(15, 2)),
+    ("P_COMMENT", pa.string()),
 ]
 
 all_schemas["partsupp"] = [
-    ("PS_PARTKEY", pyarrow.int64()),
-    ("PS_SUPPKEY", pyarrow.int64()),
-    ("PS_AVAILQTY", pyarrow.int32()),
-    ("PS_SUPPLYCOST", pyarrow.decimal128(15, 2)),
-    ("PS_COMMENT", pyarrow.string()),
+    ("PS_PARTKEY", pa.int64()),
+    ("PS_SUPPKEY", pa.int64()),
+    ("PS_AVAILQTY", pa.int32()),
+    ("PS_SUPPLYCOST", pa.decimal128(15, 2)),
+    ("PS_COMMENT", pa.string()),
 ]
 
 all_schemas["region"] = [
-    ("r_REGIONKEY", pyarrow.int64()),
-    ("r_NAME", pyarrow.string()),
-    ("r_COMMENT", pyarrow.string()),
+    ("r_REGIONKEY", pa.int64()),
+    ("r_NAME", pa.string()),
+    ("r_COMMENT", pa.string()),
 ]
 
 all_schemas["supplier"] = [
-    ("S_SUPPKEY", pyarrow.int64()),
-    ("S_NAME", pyarrow.string()),
-    ("S_ADDRESS", pyarrow.string()),
-    ("S_NATIONKEY", pyarrow.int32()),
-    ("S_PHONE", pyarrow.string()),
-    ("S_ACCTBAL", pyarrow.decimal128(15, 2)),
-    ("S_COMMENT", pyarrow.string()),
+    ("S_SUPPKEY", pa.int64()),
+    ("S_NAME", pa.string()),
+    ("S_ADDRESS", pa.string()),
+    ("S_NATIONKEY", pa.int32()),
+    ("S_PHONE", pa.string()),
+    ("S_ACCTBAL", pa.decimal128(15, 2)),
+    ("S_COMMENT", pa.string()),
 ]
 
 curr_dir = os.path.dirname(os.path.abspath(__file__))
@@ -125,12 +125,12 @@
     # in to handle the trailing | in the file
     output_cols = [r[0] for r in curr_schema]
 
-    curr_schema = [pyarrow.field(r[0], r[1], nullable=False) for r in curr_schema]
+    curr_schema = [pa.field(r[0], r[1], nullable=False) for r in curr_schema]
 
     # Trailing | requires extra field for in processing
-    curr_schema.append(("some_null", pyarrow.null()))
+    curr_schema.append(("some_null", pa.null()))
 
-    schema = pyarrow.schema(curr_schema)
+    schema = pa.schema(curr_schema)
 
     source_file = os.path.abspath(
         os.path.join(curr_dir, f"../../benchmarks/tpch/data/{filename}.csv")

python/datafusion/catalog.py

@@ -24,7 +24,7 @@
 import datafusion._internal as df_internal
 
 if TYPE_CHECKING:
-    import pyarrow
+    import pyarrow as pa
 
 
 class Catalog:
@@ -67,7 +67,7 @@ def __init__(self, table: df_internal.Table) -> None:
         self.table = table
 
     @property
-    def schema(self) -> pyarrow.Schema:
+    def schema(self) -> pa.Schema:
         """Returns the schema associated with this table."""
         return self.table.schema