opteryx-catalog 0.4.11__py3-none-any.whl → 0.4.26__py3-none-any.whl

opteryx_catalog/catalog/dataset.py

@@ -9,8 +9,7 @@ from typing import Iterable
 from typing import Optional
 
 from .manifest import ParquetManifestEntry
-from .manifest import build_parquet_manifest_entry
-from .manifest import build_parquet_manifest_minmax_entry
+from .manifest import build_parquet_manifest_entry_from_bytes
 from .metadata import DatasetMetadata
 from .metadata import Snapshot
 from .metastore import Dataset
@@ -381,16 +380,20 @@ class SimpleDataset(Dataset):
         import pyarrow as pa
         import pyarrow.parquet as pq
 
+        from ..iops.fileio import WRITE_PARQUET_OPTIONS
+
         buf = pa.BufferOutputStream()
-        pq.write_table(table, buf, compression="zstd")
+        pq.write_table(table, buf, **WRITE_PARQUET_OPTIONS)
         pdata = buf.getvalue().to_pybytes()
 
         out = self.io.new_output(data_path).create()
         out.write(pdata)
         out.close()
 
-        # Build manifest entry with statistics
-        manifest_entry = build_parquet_manifest_entry(table, data_path, len(pdata))
+        # Build manifest entry with statistics using a bytes-based, per-column scan
+        manifest_entry = build_parquet_manifest_entry_from_bytes(
+            pdata, data_path, len(pdata), orig_table=table
+        )
         return manifest_entry
 
     def overwrite(self, table: Any, author: str = None, commit_message: Optional[str] = None):
@@ -542,9 +545,7 @@ class SimpleDataset(Dataset):
         }
 
         # Build new entries for files that don't already exist. Only accept
-        # Parquet files and attempt to read lightweight metadata (bytes,
-        # row count, per-column min/max) from the Parquet footer when
-        # available.
+        # Parquet files and compute full statistics for each file.
         new_entries = []
         seen = set()
         for fp in files:
@@ -555,15 +556,19 @@ class SimpleDataset(Dataset):
                 continue
             seen.add(fp)
 
-            # Attempt to read file bytes and parquet metadata
-            # Use rugo's metadata reader which is much faster (microseconds per file)
+            # Read file and compute full statistics
             try:
+                import pyarrow as pa
+                import pyarrow.parquet as pq
+
                 inp = self.io.new_input(fp)
                 with inp.open() as f:
                     data = f.read()
 
                 if data:
-                    manifest_entry = build_parquet_manifest_minmax_entry(data, fp)
+                    # Compute statistics using a single read of the compressed bytes
+                    file_size = len(data)
+                    manifest_entry = build_parquet_manifest_entry_from_bytes(data, fp, file_size)
                 else:
                     # Empty file, create placeholder entry
                     manifest_entry = ParquetManifestEntry(
@@ -581,7 +586,7 @@ class SimpleDataset(Dataset):
                         max_values=[],
                     )
             except Exception:
-                # If metadata read fails, fall back to placeholders
+                # If read fails, fall back to placeholders
                 manifest_entry = ParquetManifestEntry(
                     file_path=fp,
                     file_format="parquet",
@@ -612,9 +617,10 @@ class SimpleDataset(Dataset):
         added_files_size = 0
         added_data_size = 0
         added_records = 0
-        # Sum uncompressed sizes from new entries
+        # Sum statistics from new entries
         for entry in new_entries:
             added_data_size += entry.get("uncompressed_size_in_bytes", 0)
+            added_records += entry.get("record_count", 0)
         deleted_data_files = 0
         deleted_files_size = 0
         deleted_data_size = 0
@@ -711,7 +717,7 @@ class SimpleDataset(Dataset):
                 prev_total_records = 0
 
         # Build unique new entries (ignore duplicates in input). Only accept
-        # parquet files and try to read lightweight metadata from each file.
+        # parquet files and compute full statistics for each file.
         new_entries = []
         seen = set()
         for fp in files:
@@ -721,14 +727,7 @@ class SimpleDataset(Dataset):
                 continue
             seen.add(fp)
 
-            file_size = 0
-            record_count = 0
-            min_values = []
-            max_values = []
             try:
-                import pyarrow as pa
-                import pyarrow.parquet as pq
-
                 data = None
                 if self.io and hasattr(self.io, "new_input"):
                     inp = self.io.new_input(fp)
@@ -748,89 +747,41 @@ class SimpleDataset(Dataset):
                         data = blob.download_as_bytes()
 
                 if data:
+                    # Compute statistics using a single read of the compressed bytes
                     file_size = len(data)
-                    pf = pq.ParquetFile(pa.BufferReader(data))
-                    record_count = int(pf.metadata.num_rows or 0)
-
-                    ncols = pf.metadata.num_columns
-                    mins = [None] * ncols
-                    maxs = [None] * ncols
-                    null_counts = [0] * ncols
-                    for rg in range(pf.num_row_groups):
-                        for ci in range(ncols):
-                            col_meta = pf.metadata.row_group(rg).column(ci)
-                            stats = getattr(col_meta, "statistics", None)
-                            if not stats:
-                                continue
-                            smin = getattr(stats, "min", None)
-                            smax = getattr(stats, "max", None)
-                            snull_count = getattr(stats, "null_count", None)
-                            if smin is None and smax is None and snull_count is None:
-                                continue
-
-                            def _to_py(v):
-                                try:
-                                    return int(v)
-                                except Exception:
-                                    try:
-                                        return float(v)
-                                    except Exception:
-                                        try:
-                                            if isinstance(v, (bytes, bytearray)):
-                                                return v.decode("utf-8", errors="ignore")
-                                        except Exception:
-                                            pass
-                                        return v
-
-                            if smin is not None:
-                                sval = _to_py(smin)
-                                if mins[ci] is None:
-                                    mins[ci] = sval
-                                else:
-                                    try:
-                                        if sval < mins[ci]:
-                                            mins[ci] = sval
-                                    except Exception:
-                                        pass
-                            if smax is not None:
-                                sval = _to_py(smax)
-                                if maxs[ci] is None:
-                                    maxs[ci] = sval
-                                else:
-                                    try:
-                                        if sval > maxs[ci]:
-                                            maxs[ci] = sval
-                                    except Exception:
-                                        pass
-                            if snull_count is not None:
-                                try:
-                                    null_counts[ci] += int(snull_count)
-                                except Exception:
-                                    pass
-
-                    min_values = [m for m in mins if m is not None]
-                    max_values = [m for m in maxs if m is not None]
+                    manifest_entry = build_parquet_manifest_entry_from_bytes(data, fp, file_size)
+                else:
+                    # Empty file, create placeholder entry
+                    manifest_entry = ParquetManifestEntry(
+                        file_path=fp,
+                        file_format="parquet",
+                        record_count=0,
+                        null_counts=[],
+                        file_size_in_bytes=0,
+                        uncompressed_size_in_bytes=0,
+                        column_uncompressed_sizes_in_bytes=[],
+                        min_k_hashes=[],
+                        histogram_counts=[],
+                        histogram_bins=0,
+                        min_values=[],
+                        max_values=[],
+                    )
             except Exception:
-                file_size = 0
-                record_count = 0
-                min_values = []
-                max_values = []
-                null_counts = []
-
-            manifest_entry = ParquetManifestEntry(
-                file_path=fp,
-                file_format="parquet",
-                record_count=int(record_count),
-                null_counts=null_counts,
-                file_size_in_bytes=int(file_size),
-                uncompressed_size_in_bytes=int(file_size),  # Use compressed size as estimate
-                column_uncompressed_sizes_in_bytes=[],
-                min_k_hashes=[],
-                histogram_counts=[],
-                histogram_bins=0,
-                min_values=min_values,
-                max_values=max_values,
-            )
+                # If read fails, create placeholder entry
+                manifest_entry = ParquetManifestEntry(
+                    file_path=fp,
+                    file_format="parquet",
+                    record_count=0,
+                    null_counts=[],
+                    file_size_in_bytes=0,
+                    uncompressed_size_in_bytes=0,
+                    column_uncompressed_sizes_in_bytes=[],
+                    min_k_hashes=[],
+                    histogram_counts=[],
+                    histogram_bins=0,
+                    min_values=[],
+                    max_values=[],
+                )
             new_entries.append(manifest_entry.to_dict())
 
         manifest_path = None
@@ -850,10 +801,11 @@ class SimpleDataset(Dataset):
         added_data_files = len(new_entries)
         added_files_size = 0
         added_data_size = 0
-        # Sum uncompressed sizes from new entries
+        added_records = 0
+        # Sum statistics from new entries
         for entry in new_entries:
             added_data_size += entry.get("uncompressed_size_in_bytes", 0)
-        added_records = 0
+            added_records += entry.get("record_count", 0)
 
         total_data_files = added_data_files
         total_files_size = added_files_size
@@ -909,9 +861,7 @@ class SimpleDataset(Dataset):
         if self.catalog and hasattr(self.catalog, "save_dataset_metadata"):
             self.catalog.save_dataset_metadata(self.identifier, self.metadata)
 
-    def scan(
-        self, row_filter=None, snapshot_id: Optional[int] = None
-    ) -> Iterable[Datafile]:
+    def scan(self, row_filter=None, snapshot_id: Optional[int] = None) -> Iterable[Datafile]:
         """Return Datafile objects for the given snapshot.
 
         - If `snapshot_id` is None, use the current snapshot.
@@ -947,6 +897,393 @@ class SimpleDataset(Dataset):
         except Exception:
             return iter(())
 
+    def describe(self, snapshot_id: Optional[int] = None, bins: int = 10) -> dict:
+        """Describe all schema columns for the given snapshot.
+
+        Returns a dict mapping column name -> statistics (same shape as
+        the previous `describe` per-column output).
+        """
+        import heapq
+
+        snap = self.snapshot(snapshot_id)
+        if snap is None or not getattr(snap, "manifest_list", None):
+            raise ValueError("No manifest available for this dataset/snapshot")
+
+        manifest_path = snap.manifest_list
+
+        # Read manifest once
+        try:
+            import pyarrow as pa
+            import pyarrow.parquet as pq
+
+            inp = self.io.new_input(manifest_path)
+            with inp.open() as f:
+                data = f.read()
+
+            if not data:
+                raise ValueError("Empty manifest data")
+
+            table = pq.read_table(pa.BufferReader(data))
+            entries = table.to_pylist()
+        except Exception:
+            raise
+
+        # Resolve schema and describe all columns
+        orso_schema = None
+        try:
+            orso_schema = self.schema()
+        except Exception:
+            orso_schema = None
+
+        if orso_schema is None:
+            raise ValueError("Schema unavailable; cannot describe all columns")
+
+        # Map column name -> index for every schema column
+        col_to_idx: dict[str, int] = {c.name: i for i, c in enumerate(orso_schema.columns)}
+
+        # Initialize accumulators per column
+        stats: dict[str, dict] = {}
+        for name in col_to_idx:
+            stats[name] = {
+                "null_count": 0,
+                "mins": [],
+                "maxs": [],
+                "hashes": set(),
+                "file_hist_infos": [],
+                "min_displays": [],
+                "max_displays": [],
+                "uncompressed_bytes": 0,
+            }
+
+        total_rows = 0
+
+        def _decode_minmax(v):
+            if v is None:
+                return None
+            if isinstance(v, (int, float)):
+                return v
+            # For strings stored as string values (not bytes), return as-is
+            if isinstance(v, str):
+                # Try to parse as number for backward compatibility
+                try:
+                    return int(v)
+                except Exception:
+                    try:
+                        return float(v)
+                    except Exception:
+                        # Not a number, return the string itself for display
+                        return v
+            try:
+                if isinstance(v, (bytes, bytearray, memoryview)):
+                    b = bytes(v)
+                    if b and b[-1] == 0xFF:
+                        b = b[:-1]
+                    s = b.decode("utf-8")
+                    try:
+                        return int(s)
+                    except Exception:
+                        try:
+                            return float(s)
+                        except Exception:
+                            # Decoded bytes that aren't numbers, return as string
+                            return s
+            except Exception:
+                pass
+            return None
+
+        # Single pass through entries updating per-column accumulators
+        for ent in entries:
+            if not isinstance(ent, dict):
+                continue
+            total_rows += int(ent.get("record_count") or 0)
+
+            # prefetch lists
+            ncounts = ent.get("null_counts") or []
+            mks = ent.get("min_k_hashes") or []
+            hists = ent.get("histogram_counts") or []
+            mv = ent.get("min_values") or []
+            xv = ent.get("max_values") or []
+            mv_disp = ent.get("min_values_display") or []
+            xv_disp = ent.get("max_values_display") or []
+            col_sizes = ent.get("column_uncompressed_sizes_in_bytes") or []
+
+            for cname, cidx in col_to_idx.items():
+                # nulls
+                try:
+                    stats[cname]["null_count"] += int((ncounts or [0])[cidx])
+                except Exception:
+                    pass
+
+                # mins/maxs
+                try:
+                    raw_min = mv[cidx]
+                except Exception:
+                    raw_min = None
+                try:
+                    raw_max = xv[cidx]
+                except Exception:
+                    raw_max = None
+                dmin = _decode_minmax(raw_min)
+                dmax = _decode_minmax(raw_max)
+                if dmin is not None:
+                    stats[cname]["mins"].append(dmin)
+                if dmax is not None:
+                    stats[cname]["maxs"].append(dmax)
+
+                # collect textual display values when present
+                try:
+                    try:
+                        raw_min_disp = mv_disp[cidx]
+                    except Exception:
+                        raw_min_disp = None
+                    try:
+                        raw_max_disp = xv_disp[cidx]
+                    except Exception:
+                        raw_max_disp = None
+
+                    def _decode_display(v):
+                        if v is None:
+                            return None
+                        try:
+                            if isinstance(v, (bytes, bytearray, memoryview)):
+                                b = bytes(v)
+                                if b and b[-1] == 0xFF:
+                                    b = b[:-1]
+                                return b.decode("utf-8", errors="replace")
+                            if isinstance(v, str):
+                                return v
+                        except Exception:
+                            return None
+                        return None
+
+                    md = _decode_display(raw_min_disp)
+                    xd = _decode_display(raw_max_disp)
+                    if md is not None:
+                        stats[cname]["min_displays"].append(md)
+                    if xd is not None:
+                        stats[cname]["max_displays"].append(xd)
+                except Exception:
+                    pass
+
+                # min-k hashes
+                try:
+                    col_mk = mks[cidx] or []
+                except Exception:
+                    col_mk = []
+                for h in col_mk:
+                    try:
+                        stats[cname]["hashes"].add(int(h))
+                    except Exception:
+                        pass
+
+                # histograms
+                try:
+                    col_hist = hists[cidx]
+                except Exception:
+                    col_hist = []
+                if col_hist:
+                    try:
+                        if dmin is not None and dmax is not None and dmin != dmax:
+                            stats[cname]["file_hist_infos"].append(
+                                (float(dmin), float(dmax), list(col_hist))
+                            )
+                    except Exception:
+                        pass
+
+                # uncompressed bytes for this column (sum across files)
+                try:
+                    stats[cname]["uncompressed_bytes"] += int((col_sizes or [0])[cidx])
+                except Exception:
+                    pass
+
+        # Build results per column
+        results: dict[str, dict] = {}
+        for cname, cidx in col_to_idx.items():
+            s = stats[cname]
+            # Handle mixed types: separate strings from numbers
+            mins_filtered = [v for v in s["mins"] if v is not None]
+            maxs_filtered = [v for v in s["maxs"] if v is not None]
+
+            # Group by type: strings vs numbers
+            str_mins = [v for v in mins_filtered if isinstance(v, str)]
+            num_mins = [v for v in mins_filtered if not isinstance(v, str)]
+            str_maxs = [v for v in maxs_filtered if isinstance(v, str)]
+            num_maxs = [v for v in maxs_filtered if not isinstance(v, str)]
+
+            # Use whichever type has values (strings take precedence for text columns)
+            global_min = None
+            global_max = None
+            if str_mins:
+                global_min = min(str_mins)
+            elif num_mins:
+                global_min = min(num_mins)
+
+            if str_maxs:
+                global_max = max(str_maxs)
+            elif num_maxs:
+                global_max = max(num_maxs)
+
+            # kmv approx
+            cardinality = 0
+            cardinality_is_exact = False
+            try:
+                collected = s["hashes"]
+                if collected:
+                    smallest = heapq.nsmallest(32, collected)
+                    k = len(smallest)
+                    if k < 31:
+                        cardinality = len(set(smallest))
+                        cardinality_is_exact = True
+                    else:
+                        MAX_HASH = (1 << 64) - 1
+                        R = max(smallest)
+                        if R == 0:
+                            cardinality = len(set(smallest))
+                        else:
+                            cardinality = int((k - 1) * (MAX_HASH + 1) / (R + 1))
+            except Exception:
+                cardinality = 0
+
+            # distribution via distogram
+            distribution = None
+            if (
+                s["file_hist_infos"]
+                and global_min is not None
+                and global_max is not None
+                and global_max > global_min
+            ):
+                try:
+                    from opteryx_catalog.maki_nage.distogram import Distogram
+                    from opteryx_catalog.maki_nage.distogram import count as _count_dist
+                    from opteryx_catalog.maki_nage.distogram import count_up_to as _count_up_to
+                    from opteryx_catalog.maki_nage.distogram import merge as _merge_distogram
+                    from opteryx_catalog.maki_nage.distogram import update as _update_distogram
+
+                    dist_bin_count = max(50, bins * 5)
+                    global_d = Distogram(bin_count=dist_bin_count)
+                    for fmin, fmax, counts in s["file_hist_infos"]:
+                        fbins = len(counts)
+                        if fbins <= 0:
+                            continue
+                        temp = Distogram(bin_count=dist_bin_count)
+                        span = float(fmax - fmin) if fmax != fmin else 0.0
+                        for bi, cnt in enumerate(counts):
+                            if cnt <= 0:
+                                continue
+                            if span == 0.0:
+                                rep = float(fmin)
+                            else:
+                                rep = fmin + (bi + 0.5) * span / fbins
+                            _update_distogram(temp, float(rep), int(cnt))
+                        global_d = _merge_distogram(global_d, temp)
+
+                    distribution = [0] * bins
+                    total = int(_count_dist(global_d) or 0)
+                    if total == 0:
+                        distribution = [0] * bins
+                    else:
+                        prev = 0.0
+                        gmin = float(global_min)
+                        gmax = float(global_max)
+                        for i in range(1, bins + 1):
+                            edge = gmin + (i / bins) * (gmax - gmin)
+                            cum = _count_up_to(global_d, edge) or 0.0
+                            distribution[i - 1] = int(round(cum - prev))
+                            prev = cum
+                        diff = total - sum(distribution)
+                        if diff != 0:
+                            distribution[-1] += diff
+                except Exception:
+                    distribution = [0] * bins
+                    gspan = float(global_max - global_min)
+                    for fmin, fmax, counts in s["file_hist_infos"]:
+                        fbins = len(counts)
+                        if fbins <= 0:
+                            continue
+                        for bi, cnt in enumerate(counts):
+                            if cnt <= 0:
+                                continue
+                            rep = fmin + (bi + 0.5) * (fmax - fmin) / fbins
+                            gi = int((rep - global_min) / gspan * bins)
+                            if gi < 0:
+                                gi = 0
+                            if gi >= bins:
+                                gi = bins - 1
+                            distribution[gi] += int(cnt)
+
+            res = {
+                "dataset": self.identifier,
+                "description": getattr(self.metadata, "description", None),
+                "row_count": total_rows,
+                "column": cname,
+                "min": global_min,
+                "max": global_max,
+                "null_count": s["null_count"],
+                "uncompressed_bytes": s["uncompressed_bytes"],
+                "cardinality": cardinality,
+                "cardinality_is_exact": cardinality_is_exact,
+                "distribution": distribution,
+            }
+
+            # If textual, attempt display prefixes like describe()
+            try:
+                is_text = False
+                if orso_schema is not None:
+                    col = orso_schema.columns[cidx]
+                    ctype = getattr(col, "type", None)
+                    if ctype is not None:
+                        sctype = str(ctype).lower()
+                        if "char" in sctype or "string" in sctype or "varchar" in sctype:
+                            is_text = True
+            except Exception:
+                is_text = False
+
+            if is_text:
+                # Use only textual display values collected from manifests.
+                # Decode bytes and strip truncation marker (0xFF) if present.
+                def _decode_display_raw(v):
+                    if v is None:
+                        return None
+                    try:
+                        if isinstance(v, (bytes, bytearray, memoryview)):
+                            b = bytes(v)
+                            if b and b[-1] == 0xFF:
+                                b = b[:-1]
+                            s_val = b.decode("utf-8", errors="replace")
+                            return s_val[:16]
+                        if isinstance(v, str):
+                            return v[:16]
+                    except Exception:
+                        return None
+                    return None
+
+                min_disp = None
+                max_disp = None
+                try:
+                    if s.get("min_displays"):
+                        for v in s.get("min_displays"):
+                            dv = _decode_display_raw(v)
+                            if dv:
+                                min_disp = dv
+                                break
+                    if s.get("max_displays"):
+                        for v in s.get("max_displays"):
+                            dv = _decode_display_raw(v)
+                            if dv:
+                                max_disp = dv
+                                break
+                except Exception:
+                    min_disp = None
+                    max_disp = None
+
+                if min_disp is not None or max_disp is not None:
+                    res["min_display"] = min_disp
+                    res["max_display"] = max_disp
+
+            results[cname] = res
+
+        return results
+
     def refresh_manifest(self, agent: str, author: Optional[str] = None) -> Optional[int]:
         """Refresh manifest statistics and create a new snapshot.
 
@@ -979,8 +1316,9 @@ class SimpleDataset(Dataset):
             import pyarrow as pa
             import pyarrow.parquet as pq
 
-            prev_table = pq.read_table(pa.BufferReader(prev_data))
-            prev_rows = prev_table.to_pylist()
+            # the manifest is a parquet file, read into a pyarrow Table
+            prev_manifest = pq.read_table(pa.BufferReader(prev_data))
+            prev_rows = prev_manifest.to_pylist()
         except Exception:
             prev_rows = []
 
@@ -1000,8 +1338,8 @@ class SimpleDataset(Dataset):
                 with inp.open() as f:
                     data = f.read()
                 # Full statistics including histograms and k-hashes
-                table = pq.read_table(pa.BufferReader(data))
-                manifest_entry = build_parquet_manifest_entry(table, fp, len(data))
+                file_size = len(data)
+                manifest_entry = build_parquet_manifest_entry_from_bytes(data, fp, file_size)
                 dent = manifest_entry.to_dict()
             except Exception:
                 # Fall back to original entry if re-read fails