dataiter 2.0__tar.gz → 2.1__tar.gz

{dataiter-2.0 → dataiter-2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: dataiter
-Version: 2.0
+Version: 2.1
 Summary: Simple, light-weight data frames for Python
 Project-URL: Homepage, https://github.com/otsaloma/dataiter
 Project-URL: Documentation, https://dataiter.readthedocs.io/

{dataiter-2.0 → dataiter-2.1}/dataiter/__init__.py

@@ -20,7 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-__version__ = "2.0"
+__version__ = "2.1"
 
 #: Default amount of elements to return from :meth:`Vector.head`,
 #: :meth:`Vector.sample` and :meth:`Vector.tail`.
@@ -89,12 +89,20 @@ from dataiter.aggregate import std
 from dataiter.aggregate import sum
 from dataiter.aggregate import var
 
+from dataiter.window import cummax
+from dataiter.window import cummin
+from dataiter.window import cumsum
+from dataiter.window import rank
+from dataiter.window import shift
+
 from dataiter.io import read_csv
+from dataiter.io import read_duckdb
 from dataiter.io import read_geojson
 from dataiter.io import read_json
 from dataiter.io import read_jsonl
 from dataiter.io import read_npz
 from dataiter.io import read_parquet
+from dataiter.io import read_parquet_chunked
 
 from dataiter import util
 

{dataiter-2.0 → dataiter-2.1}/dataiter/aggregate.py

@@ -49,7 +49,8 @@ def all(x, *, drop_na=True):
     Return whether all elements of `x` evaluate to ``True``.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     Equivalent to ``numpy.all``, see the NumPy documentation for details:
     https://numpy.org/doc/stable/reference/generated/numpy.all.html
@@ -86,7 +87,8 @@ def any(x, *, drop_na=True):
     Return whether any element of `x` evaluates to ``True``.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     Equivalent to ``numpy.any``, see the NumPy documentation for details:
     https://numpy.org/doc/stable/reference/generated/numpy.any.html
@@ -107,10 +109,10 @@ def count(x="", *, drop_na=False):
     Return the amount of elements in `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`. Since
-    all columns in a data frame should have the same amount of elements (i.e.
-    rows), you can just leave the x argument at its default blank string, which
-    will give you that row count.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`. Since all columns in a data frame should have the
+    same amount of elements (i.e. rows), you can just leave the x argument at
+    its default blank string, which will give you that row count.
 
     >>> di.count(di.Vector([1, 2, 3]))
     >>> di.count()
@@ -123,6 +125,8 @@ def count(x="", *, drop_na=False):
             return np.diff(starts, append=data.nrow)
         aggregate.group_aware = True
         return aggregate
+    if not isinstance(x, (Vector, str)):
+        raise TypeError("Expected Vector or str")
     return dispatch(count_groups, x)
 
 @composite
@@ -131,21 +135,20 @@ def count_unique(x, *, drop_na=False):
     Return the amount of unique elements in `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.count_unique(di.Vector([1, 2, 2, 3, 3, 3]))
     >>> di.count_unique("x")
     """
-    def count_unique_groups(x, starts):
-        y, starts = handle_na_groups(x, starts, drop_na)
-        n = np.diff(starts, append=len(y))
+    def count_unique_groups(codes, starts, take):
+        n = np.diff(starts, append=len(codes))
         group = np.repeat(np.arange(len(starts)), n)
-        codes = value_codes(y)
         m = codes.max(initial=0) + 1
         # Count unique pairs of group and value.
         keys = util.unique_codes(group * m + codes)[0]
         return np.bincount(keys // m, minlength=len(starts))
-    return dispatch(count_unique_groups, x)
+    return dispatch_codes(count_unique_groups, x, drop_na)
 
 def datetime_as_float(x):
     if not (x.is_datetime() or x.is_timedelta()): return x
@@ -171,6 +174,34 @@ def dispatch(function, x):
         return aggregate
     return item(function(x, np.array([0]))[0])
 
+def dispatch_codes(function, x, drop_na):
+    # Like dispatch, but function takes value codes of x instead of x, starts
+    # of their groups and a function returning elements of x at indices of
+    # codes. For a column, codes are computed in source order and reordered
+    # to group order, which is faster than reordering e.g. strings first.
+    # With drop_na, missing values get code -1, hashing only the rest.
+    def get_codes(x):
+        na = x.is_na() if drop_na else None
+        if na is None or not na.any():
+            return value_codes(x)
+        codes = np.full(len(x), -1, np.int64)
+        codes[~na] = value_codes(x[~na])
+        return codes
+    def call(codes, take, starts):
+        valid = codes >= 0
+        if valid.all():
+            return function(codes, starts, take)
+        kept = np.flatnonzero(valid)
+        starts = (np.cumsum(valid) - valid)[starts]
+        return function(codes[kept], starts, lambda i: take(kept[i]))
+    if isinstance(x, str):
+        def aggregate(data, starts):
+            take = functools.partial(data._peek, x)
+            return call(data._map(x, get_codes), take, starts)
+        aggregate.group_aware = True
+        return aggregate
+    return item(call(get_codes(x), x.__getitem__, np.array([0]))[0])
+
 def fill_missing(values, found):
     if found.all(): return values
     # Upcast to have missing values, e.g. int to float.
@@ -183,7 +214,8 @@ def first(x, *, drop_na=False):
     Return the first element of `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.first(di.Vector([1, 2, 3]))
     >>> di.first("x")
@@ -214,7 +246,8 @@ def last(x, *, drop_na=False):
     Return the last element of `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.last(di.Vector([1, 2, 3]))
     >>> di.last("x")
@@ -227,7 +260,8 @@ def max(x, *, drop_na=True):
     Return the maximum of elements in `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.max(di.Vector([4, 5, 6]))
     >>> di.max("x")
@@ -240,9 +274,9 @@ def mean(x, *, weights=None, drop_na=True):
     Return the arithmetic mean of `x`, optionally weighted by `weights`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`,
-    weighted by column `weights`, if given. If `x` is a vector, `weights`
-    should be a vector of the same length.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`, weighted by column `weights`, if given. If `x` is
+    a vector, `weights` should be a vector of the same length.
 
     With `drop_na`, rows where either `x` or `weights` is missing are
     dropped. If weights sum to zero, the result is NaN.
@@ -295,7 +329,8 @@ def median(x, *, drop_na=True):
     Return the median of `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     Equivalent to ``numpy.median``, see the NumPy documentation for details:
     https://numpy.org/doc/stable/reference/generated/numpy.median.html
@@ -331,7 +366,8 @@ def min(x, *, drop_na=True):
     Return the minimum of elements in `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.min(di.Vector([4, 5, 6]))
     >>> di.min("x")
@@ -364,18 +400,17 @@ def mode(x, *, drop_na=True):
     Return the most common value in `x`, the first one if tied.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.mode(di.Vector([1, 2, 2, 3, 3, 3]))
     >>> di.mode("x")
     """
-    def mode_groups(x, starts):
-        y, starts = handle_na_groups(x, starts, drop_na)
-        n = np.diff(starts, append=len(y))
+    def mode_groups(codes, starts, take):
+        n = np.diff(starts, append=len(codes))
         found = n > 0
         if found.any():
             group = np.repeat(np.arange(len(starts)), n)
-            codes = value_codes(y)
             m = codes.max() + 1
             keys, first, pairs = util.unique_codes(group * m + codes)
             # Unique pairs of group and value are in order of first appearance
@@ -388,11 +423,11 @@ def mode(x, *, drop_na=True):
             most = np.repeat(most, np.diff(pair_starts, append=len(keys)))
             candidates = np.flatnonzero(count == most)
             chosen = candidates[group_starts(pair_group[candidates])]
-            values = y[first[chosen]]
+            values = take(first[chosen])
         else:
-            values = y[:0]
+            values = take(np.array([], np.int64))
         return fill_missing(values, found)
-    return dispatch(mode_groups, x)
+    return dispatch_codes(mode_groups, x, drop_na)
 
 @composite
 def nth(x, index, *, drop_na=False):
@@ -400,7 +435,8 @@ def nth(x, index, *, drop_na=False):
     Return the element of `x` at `index` (zero-based).
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.nth(di.Vector([1, 2, 3]), 1)
     >>> di.nth("x", 1)
@@ -429,7 +465,8 @@ def prod(x, *, drop_na=True):
     Return the product of `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.prod(di.Vector([1, 2, 3]))
     >>> di.prod("x")
@@ -442,7 +479,8 @@ def quantile(x, q, *, drop_na=True):
     Return the `qth` quantile of `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     Equivalent to ``numpy.quantile``, see the NumPy documentation for details:
     https://numpy.org/doc/stable/reference/generated/numpy.quantile.html
@@ -502,7 +540,8 @@ def std(x, *, ddof=0, drop_na=True):
     Return the standard deviation of `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     Equivalent to ``numpy.std``, see the NumPy documentation for details:
     https://numpy.org/doc/stable/reference/generated/numpy.std.html
@@ -526,7 +565,8 @@ def sum(x, *, drop_na=True):
     Return the sum of `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     >>> di.sum(di.Vector([1, 2, 3]))
     >>> di.sum("x")
@@ -562,7 +602,8 @@ def var(x, *, ddof=0, drop_na=True):
     Return the variance of `x`.
 
     If `x` is a string, return a function usable with
-    :meth:`.DataFrame.aggregate` that operates group-wise on column `x`.
+    :meth:`.DataFrame.aggregate` and :meth:`.DataFrame.modify` that operates
+    group-wise on column `x`.
 
     Equivalent to ``numpy.var``, see the NumPy documentation for details:
     https://numpy.org/doc/stable/reference/generated/numpy.var.html