PySpark Column Selectors

A small object-oriented framework for selecting DataFrame columns by dtype, position, or name pattern instead of typing out column names by hand. A selector can be passed almost anywhere a column or list of columns is normally accepted, and selectors can be combined with set operations and chained with ordinary column transforms.

Adds column selector functionality to pyspark

• GitHub: https://github.com/tmichel3796/PySpark_Column_Selectors/
• PyPI package: https://pypi.org/project/PySpark_Column_Selectors/
• Created by: Trevor A. Michel https://TrevorMichel.com | GitHub https://github.com/tmichel3796 | PyPI https://pypi.org/user/tmichel3796/
• Free software: MIT License

Interactive ReadMe

To view this readme as a interactive google collab document please open the link below. https://colab.research.google.com/drive/1JRyF2u-OgNuzsaxjGJJeHjldf_znptI9?usp=sharing

Dependencies

• Python 3.8+
• PySpark 3.4 or newer (for full functionality)
  • PySpark 3.3 still works, except T.TimestampNTZType won't exist
  • temporal()/datetime_() fall back to matching only T.TimestampType on those older versions instead of raising.
  • withColumns with multiple selector-driven mutations in a single call (see Using selectors with other DataFrame operations) requires PySpark 3.3+, since DataFrame.withColumns itself didn't exist before then.
• No third-party packages beyond PySpark itself -- everything else used (functools, operator, re, typing) is part of the Python standard library.
• Spark Connect (pyspark.sql.connect) is supported automatically when present, but is not required -- this framework works the same either way.

Importing

import PySparkSelectors as scs

Every example below assumes this import, and calls each selector as scs.<function_name>(...).

Example data

Every example in this document runs against the same small DataFrame:

df = spark.createDataFrame(
    [
        (1, "north", "east", 120.5, 88, True, "2024-01-15", 910, "2024-06-01 08:30:00"),
        (2, "south", "west", 340.0, 95, False, "2024-03-22", 875, "2024-06-02 14:45:00"),
    ],
    ["id", "name", "region", "elevation", "totalscore", "isactive", "signupdate", "score2024", "lastseen"],
)

Dtypes: id (int), name/region (string), elevation (double), totalscore/score2024 (int), isactive (boolean), signupdate (date), lastseen (timestamp).

Dtype-based selectors

Dtype selectors match against the real pyspark.sql.types classes, not dtype strings. Import the types module alongside the selectors:

import pyspark.sql.types as T

by_dtype([...])

Matches any pyspark.sql.types.DataType class (or classes) you pass in. A category base class (e.g. T.NumericType) matches every concrete subclass of it too -- by_dtype([T.NumericType]) matches integer, floating-point, and decimal columns in one call.

df.select(scs.by_dtype([T.StringType]))

string()

df.select(scs.string())

boolean()

df.select(scs.boolean())

binary()

There are no binary columns in this DataFrame, so this raises an error by default rather than returning nothing -- see Handling zero matches below.

integer()

Fixed-width integer columns (tinyint/smallint/int/bigint).

df.select(scs.integer())

floats()

Floating-point columns (float/double).

df.select(scs.floats())

numeric()

Every fixed-width numeric column (integer or floating-point).

df.select(scs.numeric())

date()

df.select(scs.date())

datetime_()

df.select(scs.datetime_())

temporal()

Date or timestamp columns.

df.select(scs.temporal())

Position-based selectors

by_index(*indexes)

Selects one or more columns by position. Negative positions work the same way Python indexing does (-1 is the last column), and a range(...) can be mixed in with plain integers to grab a span in one call.

df.select(scs.by_index(range(0, 1), 3))

first()

df.select(scs.first())

last()

df.select(scs.last())

Handle column not in range error {#handle-column-not-in-range-error}

By default, by_index raises if any requested position is out of range for the DataFrame (this DataFrame has 9 columns, at positions 0-8):

df.select(scs.by_index(9))
# raises pyspark.errors.PySparkValueError -- position 9 is out of range

Passing strict_index_bounds=False skips any out-of-range position instead of raising, so only the in-range positions are matched:

df.select(scs.by_index(0, 9, strict_index_bounds=False))

Note: if every requested position ends up out of range, strict_index_bounds=False alone still isn't enough to avoid an error -- the selector would then match zero columns, which raises on its own (see Handling zero matches). To allow that case too, also pass require_col_match=False.

Name-based selectors

All name-based matching below is case-insensitive -- scs.by_name("REGION") matches a column named region just as well as scs.by_name("region") does, and column names may contain a literal dot (e.g. meta.source) without being mistaken for a nested field.

matches(pattern)

Selects columns whose name matches a regular expression.

df.select(scs.matches("score"))

matches takes any pattern the re module can compile, not just a plain substring -- anchors, character classes, and alternation all work:

df.select(scs.matches(r"^s.*[aeiou]$|\d"))

This pattern matches a column name if it either starts with s and ends in a vowel (signupdate), or contains a digit anywhere in the name (score2024).

starts_with(*prefixes)

df.select(scs.starts_with("e"))

ends_with(*suffixes)

df.select(scs.ends_with("date"))

contains(*substrings)

df.select(scs.contains("2024"))

by_name(*names)

Selects an exact list of column names.

df.select(scs.by_name("id", "region"))

exclude(*names)

Selects every column except the given names.

df.select(scs.exclude("region"))

alpha()

Column names made up only of letters. score2024 is the only column name containing a digit, so it's the only one excluded here.

df.select(scs.alpha())

alphanumeric()

Column names made up only of letters and/or digits -- every column name in this DataFrame qualifies, including score2024.

df.select(scs.alphanumeric())

Returns every column, identical to the example data table above.

all()

df.select(scs.all())

Returns every column, identical to the example data table above.

Combining selectors

Selectors support set-style operators so more specific selections can be built without listing columns by hand. Combined selectors always resolve columns in the same order they appear in the underlying DataFrame.

Union (|)

df.select(scs.string() | scs.starts_with("e"))

Intersection (&)

df.select(scs.numeric() & scs.matches("score"))

Difference (-)

df.select(scs.all() - scs.numeric())

Symmetric difference (^)

Columns matched by exactly one side, not both.

df.select(scs.contains("score") ^ scs.numeric())

Complement (~)

df.select(~scs.temporal())

Chaining transforms onto a selector

A transform can be chained onto a selector before it's resolved. The transform is applied to every column the selector matches, and the original column name is preserved in the result.

.upper()

df.select(scs.string().upper())

.cast(...)

df.select(scs.floats().cast("string"))

The values look the same, but elevation is now a string column instead of a double column.

Arithmetic

df.select(scs.integer() + 1)

Because every column function available in PySpark is also available as a chained method on a selector, this works with things like .substr(...), .isNull(), .between(...), and comparison operators too, not just a fixed list of built-ins.

Renaming columns:

.prefix(...), .suffix(...), and .map_alias(...) Three functions used to modify a columns name.

.prefix(prefix)

Prepends prefix to every matched column's name.

df.select(scs.by_name("id", "region").prefix("src_"))

.suffix(suffix)

Appends suffix to every matched column's name.

df.select(scs.numeric().suffix("_raw"))

.map_alias(func)

Renames every matched column by passing its current name to func, a one-argument function returning the new name. .prefix(...) and .suffix(...) are both just .map_alias(...) with a prepend/append built in -- .map_alias(...) is for anything more custom.

df.select(scs.by_name("name", "region").map_alias(lambda n: n.upper()))

Note this renames the columns, not the values -- compare to .upper() above, which uppercases the values but keeps the column names unchanged.

Using selectors with row filtering

A selector used inside a comparison becomes a row condition instead of a column pick, and conditions built this way can be combined with &, |, and ~ the same way plain columns can. filter/where still return every original column for the matching rows -- only the relevant column(s) are shown below for readability.

A single condition

df.filter(scs.by_name("totalscore") > 90)

(the row for id=1 is dropped, since totalscore there is 88)

Combining conditions with &

df.filter((scs.by_name("totalscore") > 90) & (scs.by_name("elevation") > 200))

Using selectors with other DataFrame operations {#using-selectors-with-other-dataframe-operations}

select with multiple selectors

df.select(scs.string(), scs.integer())

groupBy / agg

df.groupBy(scs.by_name("region")).agg(scs.integer().sum())

drop

df.drop(scs.by_name("lastseen", "signupdate"))

Removes lastseen and signupdate, leaving id, name, region, elevation, totalscore, isactive, score2024.

sort / orderBy

df.sort(scs.by_name("elevation").desc())

(every other original column is still present in the result; only elevation is shown here for readability)

withColumns with multiple selector-driven mutations

df.withColumns(
    scs.floats().cast("string"),
    scs.string().upper(),
)

Both mutations happen in the same call: every double column is cast to string, and every string column is uppercased.

Checking whether something is a selector

scs.is_selector(scs.numeric())  # True
scs.is_selector("totalscore")  # False

Handling zero matches {#handling-zero-matches}

By default, if a selector resolves to zero columns against a given DataFrame, it raises an error rather than silently returning nothing. This also applies to combined selectors -- for example, subtracting a selector that covers every matching column from itself will raise, since the result has nothing left in it.

df.select(scs.binary())
# raises pyspark.errors.PySparkValueError -- zero columns matched

Every selector function accepts a require_col_match keyword argument, which defaults to True. Passing require_col_match=False opts out of this check for that particular selector, allowing a zero-column match to pass through silently instead of raising:

df.select(scs.binary(require_col_match=False))
# returns a result with zero columns instead of raising

This also applies to by_dtype, by_index, and matches directly:

df.select(scs.all() - scs.by_index(range(0, 8)))
# There are 9 total columns in the df.
# As such the above column selector expression returns all columns - all columns.
# as such no columns of data are returned.
# Result: raises pyspark.errors.PySparkValueError -- zero columns matched

To bypass the no columns returned error, see the below examples. All column selector functions have access to require_col_match by_index is the only one that offers access to strict_index_bounds as shown in Handle column not in range error

df.select(scs.by_dtype([T.BinaryType], require_col_match=False))
df.select(scs.by_index(9, strict_index_bounds=False, require_col_match=False))
df.select(scs.matches("no_such_pattern", require_col_match=False))

Author

PySparkSelectors was created in 2026 by Trevor A. Michel.

Built with Cookiecutter and the audreyfeldroy/cookiecutter-pypackage project template.